vpd-permission-poc/ai-web-agent-console/ai_web_agent_console/mcp_tool_router.py

"""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",
]