import asyncio
from raganything import RAGAnything
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
from lightrag.utils import EmbeddingFunc

async def main():
    # Initialize RAGAnything
    rag = RAGAnything(
        llm_model_func=lambda prompt, system_prompt=None, history_messages=[], **kwargs: openai_complete_if_cache(
            "deepseek-chat",
            prompt,
            system_prompt=system_prompt,
            history_messages=history_messages,
            api_key="sk-44ae895eeb614aa1a9c6460579e322f1",
            base_url="https://api.deepseek.com"
            **kwargs,
        ),
        vision_model_func=lambda prompt, system_prompt=None, history_messages=[], image_data=None, **kwargs: openai_complete_if_cache(
            "GLM-4.1V-9B-Thinking",
            "",
            system_prompt=None,
            history_messages=[],
            messages=[
                {"role": "system", "content": system_prompt} if system_prompt else None,
                {"role": "user", "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}}
                ]} if image_data else {"role": "user", "content": prompt}
            ],
            api_key="sk-pbqibyjwhrgmnlsmdygplahextfaclgnedetybccknxojlyl",
            base_url='https://api.siliconflow.cn/v1/chat/completions'
            **kwargs,
        ) if image_data else openai_complete_if_cache(
            "GLM-4.1V-9B-Thinking",
            prompt,
            system_prompt=system_prompt,
            history_messages=history_messages,
            api_key="sk-pbqibyjwhrgmnlsmdygplahextfaclgnedetybccknxojlyl",
            base_url='https://api.siliconflow.cn/v1/chat/completions'
            **kwargs,
        ),
        embedding_func=EmbeddingFunc(
            embedding_dim=3072,
            max_token_size=8192,
            func=lambda texts: openai_embed(
                texts,
                model="BAAI/bge-m3",
                api_key="sk-pbqibyjwhrgmnlsmdygplahextfaclgnedetybccknxojlyl",
                base_url="https://api.siliconflow.cn/v1/embeddings",
            ),
        ),
    )

    # Process a document
    await rag.process_document_complete(
        file_path="../Txt/黄琬乔2023蓝桥杯省赛准考证.pdf",
        output_dir="./output",
        parse_method="auto"
    )

    # Query the processed content
    # Pure text query - for basic knowledge base search
    text_result = await rag.aquery(
        "这篇文档中说了什么内容?",
        mode="hybrid"
    )
    print("Text query result:", text_result)

#     # Multimodal query with specific multimodal content
#     multimodal_result = await rag.aquery_with_multimodal(
#     "Explain this formula and its relevance to the document content",
#     multimodal_content=[{
#         "type": "equation",
#         "latex": "P(d|q) = \\frac{P(q|d) \\cdot P(d)}{P(q)}",
#         "equation_caption": "Document relevance probability"
#     }],
#     mode="hybrid"
# )
#     print("Multimodal query result:", multimodal_result)

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