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())