router_embedding.py

from typing import Callable, List, Optional, TYPE_CHECKING

from numpy import mean

from mcp_agent.agents.agent import Agent
from mcp_agent.workflows.embedding.embedding_base import (
    EmbeddingModel,
    FloatArray,
    compute_similarity_scores,
    compute_confidence,
)
from mcp_agent.workflows.router.router_base import (
    Router,
    RouterCategory,
    RouterResult,
)

if TYPE_CHECKING:
    from mcp_agent.context import Context


class EmbeddingRouterCategory(RouterCategory):
    """A category for embedding-based routing"""

    embedding: FloatArray | None = None
    """Pre-computed embedding for this category"""


class EmbeddingRouter(Router):
    """
    A router that uses embedding similarity to route requests to appropriate categories.
    This class helps to route an input to a specific MCP server, an Agent (an aggregation of MCP servers),
    or a function (any Callable).

    Features:
    - Semantic similarity based routing using embeddings
    - Flexible embedding model support
    - Support for formatting and combining category metadata

    Example usage:
        # Initialize router with embedding model
        router = EmbeddingRouter(
            embedding_model=OpenAIEmbeddingModel(model="text-embedding-3-small"),
            mcp_servers_names=["customer_service", "tech_support"],
        )

        # Route a request
        results = await router.route("My laptop keeps crashing")
    """

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

        self.embedding_model = embedding_model

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

    async def initialize(self):
        """Initialize by computing embeddings for all categories"""

        async def create_category_with_embedding(
            category: RouterCategory,
        ) -> EmbeddingRouterCategory:
            # Get formatted text representation of category
            category_text = self.format_category(category)
            embedding = self._compute_embedding([category_text])
            category_with_embedding = EmbeddingRouterCategory(
                **category, embedding=embedding
            )

            return category_with_embedding

        if self.initialized:
            return

        # Create categories for servers, agents, and functions
        await super().initialize()
        self.initialized = False  # We are not initialized yet

        for name, category in self.server_categories.items():
            category_with_embedding = await create_category_with_embedding(category)
            self.server_categories[name] = category_with_embedding
            self.categories[name] = category_with_embedding

        for name, category in self.agent_categories.items():
            category_with_embedding = await create_category_with_embedding(category)
            self.agent_categories[name] = category_with_embedding
            self.categories[name] = category_with_embedding

        for name, category in self.function_categories.items():
            category_with_embedding = await create_category_with_embedding(category)
            self.function_categories[name] = category_with_embedding
            self.categories[name] = category_with_embedding

        self.initialized = True

    async def route(
        self, request: str, top_k: int = 1
    ) -> List[RouterResult[str | Agent | Callable]]:
        """Route the request based on embedding similarity"""
        if not self.initialized:
            await self.initialize()

        return await self._route_with_embedding(request, top_k)

    async def route_to_server(
        self, request: str, top_k: int = 1
    ) -> List[RouterResult[str]]:
        """Route specifically to server categories"""
        if not self.initialized:
            await self.initialize()

        results = await self._route_with_embedding(
            request,
            top_k,
            include_servers=True,
            include_agents=False,
            include_functions=False,
        )
        return [r.result for r in results[:top_k]]

    async def route_to_agent(
        self, request: str, top_k: int = 1
    ) -> List[RouterResult[Agent]]:
        """Route specifically to agent categories"""
        if not self.initialized:
            await self.initialize()

        results = await self._route_with_embedding(
            request,
            top_k,
            include_servers=False,
            include_agents=True,
            include_functions=False,
        )
        return [r.result for r in results[:top_k]]

    async def route_to_function(
        self, request: str, top_k: int = 1
    ) -> List[RouterResult[Callable]]:
        """Route specifically to function categories"""
        if not self.initialized:
            await self.initialize()

        results = await self._route_with_embedding(
            request,
            top_k,
            include_servers=False,
            include_agents=False,
            include_functions=True,
        )
        return [r.result for r in results[:top_k]]

    async def _route_with_embedding(
        self,
        request: str,
        top_k: int = 1,
        include_servers: bool = True,
        include_agents: bool = True,
        include_functions: bool = True,
    ) -> List[RouterResult]:
        def create_result(category: RouterCategory, request_embedding):
            if category.embedding is None:
                return None

            similarity = compute_similarity_scores(
                request_embedding, category.embedding
            )

            return RouterResult(
                p_score=compute_confidence(similarity), result=category.category
            )

        request_embedding = self._compute_embedding([request])

        results: List[RouterResult] = []
        if include_servers:
            for _, category in self.server_categories.items():
                result = create_result(category, request_embedding)
                if result:
                    results.append(result)

        if include_agents:
            for _, category in self.agent_categories.items():
                result = create_result(category, request_embedding)
                if result:
                    results.append(result)

        if include_functions:
            for _, category in self.function_categories.items():
                result = create_result(category, request_embedding)
                if result:
                    results.append(result)

        results.sort(key=lambda x: x.p_score, reverse=True)
        return results[:top_k]

    async def _compute_embedding(self, data: List[str]):
        # Get embedding for the provided text
        embeddings = await self.embedding_model.embed(data)

        # Use mean pooling to combine embeddings
        embedding = mean(embeddings, axis=0)

        return embedding