Indexing.py

import asyncio
import os
from ylightRag.ylightrag import LightRAG
from ylightRag.yllm.openai import gpt_4o_mini_complete, openai_embed
from ylightRag.ykg.shared_storage import initialize_pipeline_status
from ylightRag.yutils import EmbeddingFunc
Document_dir = "Documents"
"""Directory where Documents for indexing for query store"""

working_dir = "Cache" 
"""Directory where cache and temporary files are stored."""

if not os.path.exists(working_dir):
    os.mkdir(working_dir)

async def main():

    if  os.path.exists(Document_dir):
        print(f"Directory {Document_dir} exist.")
        


    rag  = LightRAG(llm_model_func=gpt_4o_mini_complete,
                    embedding_func=EmbeddingFunc(
                        embedding_dim=1536,
                        max_token_size=8192,
                        func=openai_embed
                    ),
                    working_dir=working_dir)
    
    await rag.initialize_storages()
    await initialize_pipeline_status()

    # Now indexing
    def find_txt_files(root_path):
        txt_files = []
        for root, dirs, files in os.walk(root_path):
            for file in files:
                if file.endswith(".txt"):
                    txt_files.append(os.path.join(root, file))
        return txt_files
    
    WEEK_LIST = find_txt_files(Document_dir)
    for WEEK in WEEK_LIST:
        id = WEEK_LIST.index(WEEK)
        print(f"{id}/{len(WEEK_LIST)}")
        with open(WEEK,'r',encoding="utf-8") as f:
            await rag.ainsert(f.read())
        

if __name__ == "__main__":
    asyncio.run(main())