feat: add MCP tool integration
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user