devmrko
196
ai-web-agent-console/ai_web_agent_console/mcp_tool_router.py
Normal file
196
ai-web-agent-console/ai_web_agent_console/mcp_tool_router.py
Normal file
@@ -0,0 +1,196 @@
|
||||
"""LLM-based MCP tool routing.
|
||||
|
||||
The router receives only the user question and the discovered MCP tool
|
||||
descriptors. It never receives MCP bearer tokens or provider credentials.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
from typing import Any, Mapping
|
||||
|
||||
from ai_web_agent_console.oci_genai_sdk import (
|
||||
build_oci_genai_completion_client,
|
||||
temperature_for_model_profile,
|
||||
)
|
||||
from ai_web_agent_console.model_registry import resolve_model_profile
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpTool:
|
||||
name: str
|
||||
description: str
|
||||
schema: Mapping[str, Any]
|
||||
read_only: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoutedMcpTool:
|
||||
server_id: str
|
||||
tool: McpTool
|
||||
|
||||
|
||||
class McpToolRouterError(RuntimeError):
|
||||
"""Safe routing error. Must not contain secrets or provider traces."""
|
||||
|
||||
|
||||
def route_mcp_tool_across_servers_with_llm(
|
||||
tools: list[RoutedMcpTool],
|
||||
question: str,
|
||||
*,
|
||||
router_model_profile: str,
|
||||
) -> RoutedMcpTool:
|
||||
"""Select one discovered MCP server/tool pair with OCI GenAI."""
|
||||
|
||||
candidates = list(tools)
|
||||
if not candidates:
|
||||
raise McpToolRouterError("라우팅 가능한 MCP tool이 없습니다.")
|
||||
|
||||
by_key = {
|
||||
"{}::{}".format(candidate.server_id, candidate.tool.name): candidate
|
||||
for candidate in candidates
|
||||
}
|
||||
route_keys = list(by_key)
|
||||
try:
|
||||
profile = resolve_model_profile(router_model_profile)
|
||||
client = build_oci_genai_completion_client(
|
||||
profile.model_id,
|
||||
profile.answer_model_region,
|
||||
profile.answer_model_endpoint,
|
||||
)
|
||||
tool_catalog = [
|
||||
{
|
||||
"route_key": "{}::{}".format(candidate.server_id, candidate.tool.name),
|
||||
"server_id": candidate.server_id,
|
||||
"tool_name": candidate.tool.name,
|
||||
"description": candidate.tool.description[:1000],
|
||||
"input_properties": sorted(
|
||||
(
|
||||
candidate.tool.schema.get("properties", {})
|
||||
if isinstance(
|
||||
candidate.tool.schema.get("properties"), Mapping
|
||||
)
|
||||
else {}
|
||||
).keys()
|
||||
),
|
||||
"read_only": candidate.tool.read_only,
|
||||
}
|
||||
for candidate in candidates
|
||||
]
|
||||
text = client.complete(
|
||||
system_prompt=(
|
||||
"You are an MCP server and tool router. Choose exactly one "
|
||||
"server/tool route for the user question from the discovered "
|
||||
"routes. Return only JSON that matches the schema. Never "
|
||||
"request or expose bearer tokens. Do not invent server ids or "
|
||||
"tool names."
|
||||
),
|
||||
user_prompt=json.dumps(
|
||||
{
|
||||
"question": question,
|
||||
"routes": tool_catalog,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
response_schema={
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"required": ["route_key"],
|
||||
"properties": {
|
||||
"route_key": {
|
||||
"type": "string",
|
||||
"enum": route_keys,
|
||||
}
|
||||
},
|
||||
},
|
||||
max_tokens=256,
|
||||
temperature=temperature_for_model_profile(profile),
|
||||
)
|
||||
routed = json.loads(text)
|
||||
except Exception:
|
||||
raise McpToolRouterError("LLM tool router 호출에 실패했습니다.") from None
|
||||
|
||||
if not isinstance(routed, Mapping):
|
||||
raise McpToolRouterError("LLM tool router 응답 형식이 올바르지 않습니다.")
|
||||
selected_key = str(routed.get("route_key") or "").strip()
|
||||
selected = by_key.get(selected_key)
|
||||
if selected is None:
|
||||
raise McpToolRouterError("LLM tool router가 허용되지 않은 route를 선택했습니다.")
|
||||
return selected
|
||||
|
||||
|
||||
def route_mcp_tool_with_llm(
|
||||
tools: list[McpTool],
|
||||
question: str,
|
||||
*,
|
||||
preferred_tool: str,
|
||||
tool_allowlist: tuple[str, ...],
|
||||
router_model_profile: str,
|
||||
) -> McpTool:
|
||||
"""Backward-compatible single-server routing helper."""
|
||||
|
||||
candidates = [
|
||||
tool for tool in tools if not tool_allowlist or tool.name in tool_allowlist
|
||||
]
|
||||
routed = route_mcp_tool_across_servers_with_llm(
|
||||
[RoutedMcpTool(server_id="default", tool=tool) for tool in candidates],
|
||||
question,
|
||||
router_model_profile=router_model_profile,
|
||||
)
|
||||
return routed.tool
|
||||
|
||||
|
||||
def build_mcp_tool_arguments(
|
||||
tool: McpTool,
|
||||
question: str,
|
||||
limit: int,
|
||||
*,
|
||||
preferred_tool: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Build bounded tool arguments from the selected tool schema."""
|
||||
|
||||
properties = tool.schema.get("properties")
|
||||
if not isinstance(properties, Mapping):
|
||||
properties = {}
|
||||
|
||||
# A server's default/preferred tool still has to obey its discovered schema.
|
||||
# HMM tools use `query` and `term`; forcing the legacy `prompt`/`limit` shape
|
||||
# makes an otherwise valid tool fail argument validation.
|
||||
del preferred_tool
|
||||
input_name = next(
|
||||
(name for name in ("prompt", "question", "query", "term", "text") if name in properties),
|
||||
"",
|
||||
)
|
||||
if not input_name:
|
||||
required = tool.schema.get("required")
|
||||
if isinstance(required, list):
|
||||
input_name = next(
|
||||
(
|
||||
str(name)
|
||||
for name in required
|
||||
if isinstance(properties.get(str(name)), Mapping)
|
||||
and properties[str(name)].get("type") == "string"
|
||||
),
|
||||
"",
|
||||
)
|
||||
if input_name:
|
||||
args: dict[str, Any] = {input_name: question}
|
||||
if "max_evidence" in properties:
|
||||
args["max_evidence"] = min(limit, 10)
|
||||
elif "max_rows" in properties:
|
||||
args["max_rows"] = limit
|
||||
elif "limit" in properties:
|
||||
args["limit"] = limit
|
||||
return args
|
||||
return {"prompt": question, "limit": limit}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"McpTool",
|
||||
"McpToolRouterError",
|
||||
"RoutedMcpTool",
|
||||
"build_mcp_tool_arguments",
|
||||
"route_mcp_tool_across_servers_with_llm",
|
||||
"route_mcp_tool_with_llm",
|
||||
]
|
||||
Reference in New Issue
Block a user