110 lines
3.4 KiB
Python
110 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
PostgreSQL 数据库优化工程师智能体
|
|
|
|
功能:
|
|
- 使用 DeepSeek 模型进行智能分析
|
|
- 集成 PostgreSQL MCP 服务器工具
|
|
- 加载数据库知识库
|
|
- 提供数据库优化建议
|
|
"""
|
|
|
|
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 import 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()
|
|
self.knowledge_base = self._setup_knowledge_base()
|
|
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(
|
|
command="cmd",
|
|
args=[
|
|
"/c",
|
|
"npx",
|
|
"-y",
|
|
"@modelcontextprotocol/server-postgres",
|
|
self.config.db_connection,
|
|
],
|
|
env={},
|
|
)
|
|
return MCPTools(server_params=server_params)
|
|
|
|
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],
|
|
)
|
|
|
|
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这个数据库")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|