if __name__ == "__main__":
    """
    ## 背景

    我有一个`postgresql`数据库,现在需要一个专业的数据库优化工程师,来帮我看下设计是否合理,以及如何优化.

    ## 目标

    这个数据库优化工程师智能体应该具备以下能力:

    - [ ] 使用`claude3.7`作为模型(相当于拥有一个聪明的大脑,总是能够做出正确的决策)
    - [ ] 精通关系型数据库系统的理论知识以及`postgresql`的实现细节(知识库)
    - [ ] 能够使用`postgresql mcp server`来分析现有数据库设计(行为)

    """
    # 然后在notebook中加载
    from dotenv import load_dotenv

    # 加载当前目录下的.env文件
    load_dotenv()
    from agno.agent import Agent, RunResponse  # noqa
    from agno.models.deepseek import DeepSeek
    from agno.knowledge.pdf_url import PDFUrlKnowledgeBase
    from agno.knowledge.pdf import PDFKnowledgeBase
    from agno.vectordb.lancedb import LanceDb, SearchType
    from agno.embedder.openai import OpenAIEmbedder
    from agno.knowledge.combined import CombinedKnowledgeBase
    from agno.tools.mcp import MCPTools
    from mcp import StdioServerParameters

    # Create a knowledge base of PDFs from URLs
    # pdf_url_kb = PDFUrlKnowledgeBase(
    #     urls=["https://agno-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf"],
    #     # Use LanceDB as the vector database and store embeddings in the `recipes` table
    #     vector_db=LanceDb(
    #         table_name="recipes",
    #         uri="tmp/lancedb",
    #         search_type=SearchType.vector,
    #         embedder=OpenAIEmbedder(id="text-embedding-3-small"),
    #     ),
    # )
    # Create Local PDF knowledge base
    local_pdf_kb = PDFKnowledgeBase(
        path="D:\\Sources\\DONGJAK-TOOLS\\pdfs\\Database Fundamentals.pdf",
        vector_db=LanceDb(
            table_name="database_fundamentals",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )

    # Combine knowledge bases
    knowledge_base = CombinedKnowledgeBase(
        sources=[
            local_pdf_kb,
        ],
        vector_db=LanceDb(
            table_name="combined_documents",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )
    # Load the knowledge base: Comment after first run as the knowledge base is already loaded
    knowledge_base.load()

    server_params = StdioServerParameters(
        command="cmd",  # 或 "uvx"，取决于你的安装方式
        args=[
            "/c",
            "npx",
            "-y",
            "@modelcontextprotocol/server-postgres",
            "postgresql://postgres:postgres@192.168.1.7:5432/postgres",
        ],
        env={},  # 可选的环境变量
    )
    with MCPTools(server_params=server_params) as postgres_server:
        # 使用mcp_tools

        agent = Agent(
            model=DeepSeek(id="deepseek-chat"),
            markdown=True,
            knowledge=knowledge_base,
            search_knowledge=True,
            show_tool_calls=True,
            tools=[postgres_server],
        )

        # Get the response in a variable
        # run: RunResponse = agent.run("Share a 2 sentence horror story")
        # print(run.content)

        # Print the response in the terminal
        agent.print_response("看下aq这个数据库", stream=True)