#!/usr/bin/env python3
"""
PostgreSQL 数据库优化工程师智能体

功能:
- 使用 DeepSeek 模型进行智能分析
- 集成 PostgreSQL MCP 服务器工具
- 加载数据库知识库
- 提供数据库优化建议
"""

import asyncio
import os
from typing import Optional
from dataclasses import dataclass
from dotenv import load_dotenv

from agno.agent import Agent
from agno.models.deepseek import DeepSeek
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.client.stdio import stdio_client
from mcp import ClientSession, StdioServerParameters
from agno.knowledge.website import WebsiteKnowledgeBase

 

async def main():
    load_dotenv()
    """设置知识库"""
    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"),
        ),
    )
    # Create Website knowledge base
    website_kb = WebsiteKnowledgeBase(
        urls=["https://www.lucidchart.com/blog/database-design-best-practices"],
        max_links=10,
        vector_db=LanceDb(
            table_name="website_documents",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )

    knowledge_base = CombinedKnowledgeBase(
        sources=[local_pdf_kb, website_kb],
        vector_db=LanceDb(
            table_name="combined_documents",
            uri="tmp/lancedb",
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(id="text-embedding-3-small"),
        ),
    )
    knowledge_base.load()

    postgres_server_params = StdioServerParameters(
        command="cmd",
        args=[
            "/c",
            "npx",
            "-y",
            "@modelcontextprotocol/server-postgres",
            "postgresql://postgres:postgres@192.168.1.7:5432/aq",
        ],
        env={},
    )

    searxng_server_params = StdioServerParameters(
        command="cmd",
        args=[
            "/c",
            "npx",
            "-y",
            "https://github.com/ihor-sokoliuk/mcp-searxng.git",
        ],
        env={
            "SEARXNG_URL": "https://searchxng.ailoveworld.cn",
        },
    )
    # Create a client session to connect to the MCP server
    async with (
        MCPTools(server_params=postgres_server_params) as postgres_tools,
        MCPTools(server_params=searxng_server_params) as searxng_tools,
    ):
        agent = Agent(
            model=DeepSeek(id="deepseek-chat"),
            markdown=True,
            knowledge=knowledge_base,
            search_knowledge=True,
            show_tool_calls=True,
            tools=[postgres_tools, searxng_tools],
        )
        await agent.aprint_response("帮我分析一下aq.public数据库,并给出优化建议", stream=True)
        # await agent.aprint_response("阅读下 https://www.lucidchart.com/blog/database-design-best-practices 这篇文章", stream=True)


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