✨ feat(database_optimizer): 支持异步数据库分析

- 新增异步主函数以提升性能
- 重构知识库加载及智能体创建逻辑
- 添加客户端会话以连接MCP服务器

数据库优化工程师.py

@@ -9,6 +9,7 @@ PostgreSQL 数据库优化工程师智能体
 - 提供数据库优化建议
 """
 
+import asyncio
 import os
 from typing import Optional
 from dataclasses import dataclass
@@ -21,19 +22,23 @@ 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
+from mcp.client.stdio import stdio_client
+from mcp import ClientSession, StdioServerParameters
+
 
 @dataclass
 class DatabaseOptimizerConfig:
     """数据库优化器配置"""
+
     pdf_path: str = "D:\\Sources\\DONGJAK-TOOLS\\pdfs\\Database Fundamentals.pdf"
     db_connection: str = "postgresql://postgres:postgres@192.168.1.7:5432/postgres"
     model_id: str = "deepseek-chat"
     vector_db_path: str = "tmp/lancedb"
 
+
 class DatabaseOptimizer:
     """PostgreSQL 数据库优化引擎"""
-    
+
     def __init__(self, config: Optional[DatabaseOptimizerConfig] = None):
         self.config = config or DatabaseOptimizerConfig()
         self._load_environment()
@@ -41,34 +46,6 @@ class DatabaseOptimizer:
         self.postgres_tools = self._setup_postgres_tools()
         self.agent = self._create_agent()
 
-    def _load_environment(self):
-        """加载环境变量"""
-        load_dotenv()
-
-    def _setup_knowledge_base(self) -> CombinedKnowledgeBase:
-        """设置知识库"""
-        local_pdf_kb = PDFKnowledgeBase(
-            path=self.config.pdf_path,
-            vector_db=LanceDb(
-                table_name="database_fundamentals",
-                uri=self.config.vector_db_path,
-                search_type=SearchType.vector,
-                embedder=OpenAIEmbedder(id="text-embedding-3-small"),
-            ),
-        )
-        
-        knowledge_base = CombinedKnowledgeBase(
-            sources=[local_pdf_kb],
-            vector_db=LanceDb(
-                table_name="combined_documents",
-                uri=self.config.vector_db_path,
-                search_type=SearchType.vector,
-                embedder=OpenAIEmbedder(id="text-embedding-3-small"),
-            ),
-        )
-        knowledge_base.load()
-        return knowledge_base
-
     def _setup_postgres_tools(self) -> MCPTools:
         """设置 PostgreSQL 工具"""
         server_params = StdioServerParameters(
@@ -82,28 +59,71 @@ class DatabaseOptimizer:
             ],
             env={},
         )
-        return MCPTools(server_params=server_params)
+        # Create a client session to connect to the MCP server
+        with stdio_client(server_params) as (read, write):
+            with ClientSession(read, write) as session:
+                agent = create_filesystem_agent(session)
 
-    def _create_agent(self) -> Agent:
-        """创建智能体"""
-        return Agent(
-            model=DeepSeek(id=self.config.model_id),
-            markdown=True,
-            knowledge=self.knowledge_base,
-            search_knowledge=True,
-            show_tool_calls=True,
-            tools=[self.postgres_tools],
-        )
+                # Run the agent
+                agent.print_response(message, stream=True)
+        return MCPTools(server_params=server_params)
 
     def analyze_database(self, query: str, stream: bool = True):
         """分析数据库"""
         with self.postgres_tools:
             self.agent.print_response(query, stream=stream)
 
-def main():
-    """主入口函数"""
-    optimizer = DatabaseOptimizer()
-    optimizer.analyze_database("看下aq这个数据库")
+
+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"),
+        ),
+    )
+
+    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"),
+        ),
+    )
+    knowledge_base.load()
+
+    server_params = StdioServerParameters(
+        command="cmd",
+        args=[
+            "/c",
+            "npx",
+            "-y",
+            "@modelcontextprotocol/server-postgres",
+            "postgresql://postgres:postgres@192.168.1.7:5432/postgres",
+        ],
+        env={},
+    )
+    # Create a client session to connect to the MCP server
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write) as session:
+            postgres_tools = MCPTools(session=session)
+            await postgres_tools.initialize()
+            agent = Agent(
+                model=DeepSeek(id="deepseek-chat"),
+                markdown=True,
+                knowledge=knowledge_base,
+                search_knowledge=True,
+                show_tool_calls=True,
+                tools=[postgres_tools],
+            )
+            await agent.aprint_response("看下aq这个数据库", stream=True)
+
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())