feat: add MCP tool integration

This commit is contained in:
Xin Wang
2026-07-18 00:00:06 +08:00
parent bdf3d3dd9c
commit e39bb48ba8
21 changed files with 1641 additions and 62 deletions

View File

@@ -9,12 +9,14 @@ from db.models import (
AssistantModelBinding,
AssistantToolBinding,
KnowledgeBase,
McpServer,
ModelResource,
Tool,
)
from models import (
AssistantConfig,
RuntimeKnowledgeBase,
RuntimeMcpServer,
RuntimeModelResource,
RuntimeTool,
)
@@ -102,18 +104,55 @@ async def _tools_for(session: AsyncSession, assistant: Assistant) -> list[Runtim
.order_by(AssistantToolBinding.created_at, Tool.id)
)
).scalars().all()
return [
RuntimeTool(
id=tool.id,
name=tool.name,
function_name=tool.function_name,
type=tool.type,
description=tool.description,
definition=tool.definition or {},
secrets=tool.secrets or {},
)
server_ids = {
str(tool.mcp_server_id)
for tool in tools
]
if tool.type == "mcp" and tool.mcp_server_id
}
server_rows = (
(
await session.execute(
select(McpServer).where(
McpServer.id.in_(server_ids),
McpServer.status == "active",
)
)
).scalars().all()
if server_ids
else []
)
servers = {server.id: server for server in server_rows}
resolved: list[RuntimeTool] = []
for tool in tools:
mcp_server = None
if tool.type == "mcp":
server = servers.get(str(tool.mcp_server_id or ""))
if server is None:
continue
config = server.config or {}
secrets = server.secrets or {}
mcp_server = RuntimeMcpServer(
id=server.id,
name=server.name,
transport=server.transport,
url=server.url,
timeout_seconds=int(config.get("timeout_seconds") or 30),
headers=dict(config.get("headers") or {}),
secret_headers=dict(secrets.get("headers") or {}),
)
resolved.append(
RuntimeTool(
id=tool.id,
name=tool.name,
function_name=tool.function_name,
type=tool.type,
description=tool.description,
definition=tool.definition or {},
secrets=tool.secrets or {},
mcp_server=mcp_server,
)
)
return resolved
async def resolve_runtime_config(