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"

    # 创建 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,
    )
    # 查询
    user_prompt = "你是一个化学专家，对方程式非常熟悉！"
    result = await rag.aquery(
        user_prompt=user_prompt,
        query="氢气与氧气的反应方程式是什么？",
        mode="hybrid"
    )
    print("查询结果:", result)


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