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/DongHua"

    # 创建 LLM 模型自定义函数
    llm_model_func = create_llm_model_func(v_history_messages=[])
    # 创建可视模型自定义函数
    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 = "如果资料中提供了图片，需要把图片也展示出来。有图片路径的，需要转为![]() 这样的markdown格式展示。"
    #user_prompt = "如需画图使用mermaid格式，节点名称用英文或拼音，显示名称用中文",
    result = await rag.aquery(
        #user_prompt=user_prompt,
        #query="本文档讲述了什么内容？",
        query="吉林动画学院招聘岗位有哪些？",
        mode="hybrid"
    )
    print("查询结果:", result)


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