From 0455760852e338f6d3a56c3b9a25851659b2d214 Mon Sep 17 00:00:00 2001 From: cruld Date: Mon, 31 Mar 2025 01:55:21 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(database=5Foptimizer):=20?= =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=BC=82=E6=AD=A5=E6=95=B0=E6=8D=AE=E5=BA=93?= =?UTF-8?q?=E5=88=86=E6=9E=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增异步主函数以提升性能 - 重构知识库加载及智能体创建逻辑 - 添加客户端会话以连接MCP服务器 --- 数据库优化工程师.py | 112 ++++++++++++++++++++++++++------------------ 1 file changed, 66 insertions(+), 46 deletions(-) diff --git a/数据库优化工程师.py b/数据库优化工程师.py index df05f01..cf4c209 100644 --- a/数据库优化工程师.py +++ b/数据库优化工程师.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())