import asyncio

from raganything import RAGAnything, RAGAnythingConfig

from Util.RagUtil import create_llm_model_func, create_vision_model_func, create_embedding_func


async def main():

    config = RAGAnythingConfig(
        working_dir="./rag_storage",
        mineru_parse_method="auto",
        enable_image_processing=True,
        enable_table_processing=True,
        enable_equation_processing=True,
    )

    llm_model_func = create_llm_model_func()
    vision_model_func = create_vision_model_func(llm_model_func)
    embedding_func = create_embedding_func()

    rag = RAGAnything(
        config=config,
        llm_model_func=llm_model_func,
        vision_model_func=vision_model_func,
        embedding_func=embedding_func,
    )

    file_path = "./Txt/驿来特平台安全.docx"
    await rag.process_document_complete(
        file_path=file_path,
        output_dir="./output",
        parse_method="auto"
    )

    print("Processing complete.")

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