src/mcp_agent/workflows/router/router_llm.py

from typing import Callable, List, Literal, Optional, TYPE_CHECKING

from opentelemetry import trace
from pydantic import BaseModel

from mcp_agent.agents.agent import Agent
from mcp_agent.tracing.semconv import GEN_AI_REQUEST_TOP_K
from mcp_agent.tracing.telemetry import get_tracer
from mcp_agent.workflows.llm.augmented_llm import AugmentedLLM
from mcp_agent.workflows.router.router_base import ResultT, Router, RouterResult
from mcp_agent.logging.logger import get_logger

if TYPE_CHECKING:
    from mcp_agent.core.context import Context

logger = get_logger(__name__)


DEFAULT_ROUTING_INSTRUCTION = """
You are a highly accurate request router that directs incoming requests to the most appropriate category.
A category is a specialized destination, such as a Function, an MCP Server (a collection of tools/functions), or an Agent (a collection of servers).
Below are the available routing categories, each with their capabilities and descriptions:

{context}

Your task is to analyze the following request and determine the most appropriate categories from the options above. Consider:
- The specific capabilities and tools each destination offers
- How well the request matches the category's description
- Whether the request might benefit from multiple categories (up to {top_k})

Request: {request}

Respond in JSON format:
{{
    "categories": [
        {{
            "category": <category name>,
            "confidence": <high, medium or low>,
            "reasoning": <brief explanation>
        }}
    ]
}}

Only include categories that are truly relevant. You may return fewer than {top_k} if appropriate.
If none of the categories are relevant, return an empty list.
"""


class LLMRouterResult(RouterResult[ResultT]):
    """A class that represents the result of an LLMRouter.route request"""

    confidence: Literal["high", "medium", "low"]
    """The confidence level of the routing decision."""

    reasoning: str | None = None
    """
    A brief explanation of the routing decision.
    This is optional and may only be provided if the router is an LLM
    """


class StructuredResponseCategory(BaseModel):
    """A class that represents a single category returned by an LLM router"""

    category: str
    """The name of the category (i.e. MCP server, Agent or function) to route the input to."""

    confidence: Literal["high", "medium", "low"]
    """The confidence level of the routing decision."""

    reasoning: str | None = None
    """A brief explanation of the routing decision."""


class StructuredResponse(BaseModel):
    """A class that represents the structured response of an LLM router"""

    categories: List[StructuredResponseCategory]
    """A list of categories to route the input to."""


class LLMRouter(Router):
    """
    A router that uses an LLM to route an input to a specific category.
    """

    def __init__(
        self,
        llm: AugmentedLLM,
        server_names: List[str] | None = None,
        agents: List[Agent] | None = None,
        functions: List[Callable] | None = None,
        routing_instruction: str | None = None,
        context: Optional["Context"] = None,
        **kwargs,
    ):
        super().__init__(
            server_names=server_names,
            agents=agents,
            functions=functions,
            routing_instruction=routing_instruction,
            context=context,
            **kwargs,
        )

        self.llm = llm

    @classmethod
    async def create(
        cls,
        llm: AugmentedLLM,
        server_names: List[str] | None = None,
        agents: List[Agent] | None = None,
        functions: List[Callable] | None = None,
        routing_instruction: str | None = None,
        context: Optional["Context"] = None,
    ) -> "LLMRouter":
        """
        Factory method to create and initialize a router.
        Use this instead of constructor since we need async initialization.
        """
        instance = cls(
            llm=llm,
            server_names=server_names,
            agents=agents,
            functions=functions,
            routing_instruction=routing_instruction,
            context=context,
        )
        await instance.initialize()
        return instance

    async def route(
        self, request: str, top_k: int = 1
    ) -> List[LLMRouterResult[str | Agent | Callable]]:
        tracer = get_tracer(self.context)
        with tracer.start_as_current_span(f"{self.__class__.__name__}.route") as span:
            self._annotate_span_for_route_request(span, request, top_k)

            if not self.initialized:
                await self.initialize()

            res = await self._route_with_llm(request, top_k)
            self._annotate_span_for_router_result(span, res)
            return res

    async def route_to_server(
        self, request: str, top_k: int = 1
    ) -> List[LLMRouterResult[str]]:
        tracer = get_tracer(self.context)
        with tracer.start_as_current_span(
            f"{self.__class__.__name__}.route_to_server"
        ) as span:
            self._annotate_span_for_route_request(span, request, top_k)

            if not self.initialized:
                await self.initialize()

            res = await self._route_with_llm(
                request,
                top_k,
                include_servers=True,
                include_agents=False,
                include_functions=False,
            )
            self._annotate_span_for_router_result(span, res)
            return res

    async def route_to_agent(
        self, request: str, top_k: int = 1
    ) -> List[LLMRouterResult[Agent]]:
        tracer = get_tracer(self.context)
        with tracer.start_as_current_span(
            f"{self.__class__.__name__}.route_to_agent"
        ) as span:
            self._annotate_span_for_route_request(span, request, top_k)

            if not self.initialized:
                await self.initialize()

            res = await self._route_with_llm(
                request,
                top_k,
                include_servers=False,
                include_agents=True,
                include_functions=False,
            )
            self._annotate_span_for_router_result(span, res)
            return res

    async def route_to_function(
        self, request: str, top_k: int = 1
    ) -> List[LLMRouterResult[Callable]]:
        tracer = get_tracer(self.context)
        with tracer.start_as_current_span(
            f"{self.__class__.__name__}.route_to_function"
        ) as span:
            self._annotate_span_for_route_request(span, request, top_k)

            if not self.initialized:
                await self.initialize()

            res = await self._route_with_llm(
                request,
                top_k,
                include_servers=False,
                include_agents=False,
                include_functions=True,
            )
            self._annotate_span_for_router_result(span, res)
            return res

    async def _route_with_llm(
        self,
        request: str,
        top_k: int = 1,
        include_servers: bool = True,
        include_agents: bool = True,
        include_functions: bool = True,
    ) -> List[LLMRouterResult]:
        tracer = get_tracer(self.context)
        with tracer.start_as_current_span(
            f"{self.__class__.__name__}._route_with_llm"
        ) as span:
            self._annotate_span_for_route_request(span, request, top_k)

            if not self.initialized:
                await self.initialize()

            routing_instruction = (
                self.routing_instruction or DEFAULT_ROUTING_INSTRUCTION
            )

            # Generate the categories context
            context = self._generate_context(
                include_servers=include_servers,
                include_agents=include_agents,
                include_functions=include_functions,
            )

            # logger.debug(
            #     f"Requesting routing from LLM, \nrequest: {request} \ntop_k: {top_k} \nrouting_instruction: {routing_instruction} \ncontext={context}",
            #     data={"progress_action": "Routing", "agent_name": "LLM Router"},
            # )

            # Format the prompt with all the necessary information
            prompt = routing_instruction.format(
                context=context, request=request, top_k=top_k
            )

            # Get routes from LLM
            response = await self.llm.generate_structured(
                message=prompt,
                response_model=StructuredResponse,
            )

            if self.context.tracing_enabled:
                response_categories_data = {}
                for i, r in enumerate(response.categories):
                    response_categories_data[f"category.{i}.category"] = r.category
                    response_categories_data[f"category.{i}.confidence"] = r.confidence
                    if r.reasoning:
                        response_categories_data[f"category.{i}.reasoning"] = (
                            r.reasoning
                        )

                span.add_event(
                    "routing.response",
                    {
                        "prompt": prompt,
                        **response_categories_data,
                    },
                )

            # logger.debug(
            #     "Routing Response received",
            #     data={"progress_action": "Finished", "agent_name": "LLM Router"},
            # )

            # Construct the result
            if not response or not response.categories:
                return []

            result: List[LLMRouterResult] = []
            for r in response.categories:
                router_category = self.categories.get(r.category)
                if not router_category:
                    # Skip invalid categories
                    # TODO: saqadri - log or raise an error
                    continue

                result.append(
                    LLMRouterResult(
                        result=router_category.category,
                        confidence=r.confidence,
                        reasoning=r.reasoning,
                    )
                )

            self._annotate_span_for_router_result(span, result)

            return result[:top_k]

    def _annotate_span_for_route_request(
        self,
        span: trace.Span,
        request: str,
        top_k: int,
    ):
        """Annotate the span with the request and top_k."""
        if not self.context.tracing_enabled:
            return
        span.set_attribute("request", request)
        span.set_attribute(GEN_AI_REQUEST_TOP_K, top_k)
        span.set_attribute("llm", self.llm.name)
        span.set_attribute(
            "agents", [a.name for a in self.agents] if self.agents else []
        )
        span.set_attribute("servers", self.server_names or [])
        span.set_attribute(
            "functions", [f.__name__ for f in self.functions] if self.functions else []
        )

    def _annotate_span_for_router_result(
        self,
        span: trace.Span,
        result: List[LLMRouterResult],
    ):
        """Annotate the span with the router result."""
        if not self.context.tracing_enabled:
            return
        for i, res in enumerate(result):
            span.set_attribute(f"result.{i}.confidence", res.confidence)
            if res.reasoning:
                span.set_attribute(f"result.{i}.reasoning", res.reasoning)
            if res.p_score:
                span.set_attribute(f"result.{i}.p_score", res.p_score)

            result_key = f"result.{i}.result"
            if isinstance(res.result, str):
                span.set_attribute(result_key, res.result)
            elif isinstance(res.result, Agent):
                span.set_attribute(result_key, res.result.name)
            elif callable(res.result):
                span.set_attribute(result_key, res.result.__name__)

    def _generate_context(
        self,
        include_servers: bool = True,
        include_agents: bool = True,
        include_functions: bool = True,
    ) -> str:
        """Generate a formatted context list of categories."""

        context_list = []
        idx = 1

        # Format all categories
        if include_servers:
            for category in self.server_categories.values():
                context_list.append(self.format_category(category, idx))
                idx += 1

        if include_agents:
            for category in self.agent_categories.values():
                context_list.append(self.format_category(category, idx))
                idx += 1

        if include_functions:
            for category in self.function_categories.values():
                context_list.append(self.format_category(category, idx))
                idx += 1

        return "\n\n".join(context_list)