import asyncio

from lightrag import LightRAG
from lightrag.kg.shared_storage import initialize_pipeline_status
from raganything import RAGAnything

from Util.RagUtil import create_llm_model_func, create_vision_model_func, create_embedding_func


async def load_existing_lightrag():
    # 索引位置
    #WORKING_DIR = "./Topic/Chemistry"
    #WORKING_DIR = "./Topic/DongHua"
    #WORKING_DIR = "./Topic/Chinese"
    WORKING_DIR = "./Topic/Math"

    # 创建 LLM 模型自定义函数
    llm_model_func = create_llm_model_func()
    # 创建可视模型自定义函数
    vision_model_func = create_vision_model_func(llm_model_func)
    # 创建嵌入模型自定义函数
    embedding_func = create_embedding_func()
    # 声明LightRAG实例
    lightrag_instance = LightRAG(
        working_dir=WORKING_DIR,
        llm_model_func=llm_model_func,
        embedding_func=embedding_func
    )
    # 初始化
    await lightrag_instance.initialize_storages()
    await initialize_pipeline_status()
    # 创建RAGAnything实例，依托于LightRAG实例
    rag = RAGAnything(
        lightrag=lightrag_instance,
        vision_model_func=vision_model_func,
    )
    # 查询
    result = await rag.aquery(
        #query="氧化铁和硝酸的反应方程式?",
        query="文档介绍了哪些内容？",
        mode="hybrid"
    )
    print("查询结果:", result)


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