62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Annotated
|
|
|
|
from deepagents.middleware.filesystem import FilesystemMiddleware
|
|
from langchain.agents import create_agent
|
|
|
|
from src.agents.common import BaseAgent, BaseContext, load_chat_model
|
|
from src.agents.common.backends import create_agent_composite_backend
|
|
from src.agents.common.middlewares import RuntimeConfigMiddleware, save_attachments_to_fs
|
|
from src.agents.common.toolkits.mysql import get_mysql_tools
|
|
from src.services.mcp_service import get_mcp_server_names, get_tools_from_all_servers
|
|
from src.utils import logger
|
|
|
|
|
|
def _create_fs_backend(rt):
|
|
return create_agent_composite_backend(rt)
|
|
|
|
|
|
PROMPT = """你的任务是根据用户指令，使用数据库工具和图表工具生成 SQL 报告。
|
|
请先理解需求，再给出准确且高效的 SQL 查询；必要时调用图表工具并在最终回答中以 Markdown 图片形式展示结果。
|
|
"""
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class ReporterContext(BaseContext):
|
|
system_prompt: Annotated[str, {"__template_metadata__": {"kind": "prompt"}}] = field(
|
|
default=PROMPT,
|
|
metadata={"name": "系统提示词", "description": "描述 SQL 报告助手的行为"},
|
|
)
|
|
mcps: Annotated[list[str], {"__template_metadata__": {"kind": "mcps"}}] = field(
|
|
default_factory=lambda: ["mcp-server-chart"],
|
|
metadata={
|
|
"name": "MCP 服务",
|
|
"options": lambda: get_mcp_server_names(),
|
|
"description": "报告场景默认启用图表 MCP。",
|
|
},
|
|
)
|
|
|
|
|
|
class SqlReporterAgent(BaseAgent):
|
|
name = "数据报表助手"
|
|
description = "根据用户需求生成 SQL 查询并输出图表化报告。"
|
|
context_schema = ReporterContext
|
|
capabilities = ["file_upload", "files"]
|
|
|
|
async def get_graph(self, **kwargs):
|
|
context = self.context_schema.from_file(module_name=self.module_name)
|
|
all_mcp_tools = await get_tools_from_all_servers()
|
|
graph = create_agent(
|
|
model=load_chat_model(context.model),
|
|
system_prompt=context.system_prompt,
|
|
tools=get_mysql_tools(),
|
|
middleware=[
|
|
FilesystemMiddleware(backend=_create_fs_backend),
|
|
RuntimeConfigMiddleware(extra_tools=all_mcp_tools),
|
|
save_attachments_to_fs,
|
|
],
|
|
checkpointer=await self._get_checkpointer(),
|
|
)
|
|
logger.info("SqlReporterAgent graph created")
|
|
return graph
|