From 8d70a267eec8f93eb9837c84b8314e2d3569a6cc Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 20:04:32 +0000 Subject: [PATCH] perf(mcp): share semantic tool embeddings across pods via redis Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/AGENTS.md | 1 + .../mcp_server/semantic_tool_filter.py | 251 +++++++++++---- .../mcp_server/semantic_tool_index_store.py | 172 +++++++++++ .../proxy/hooks/mcp_semantic_filter/hook.py | 64 +++- litellm/proxy/proxy_server.py | 7 + .../mcp_server/test_semantic_tool_filter.py | 106 +++++++ .../test_semantic_tool_index_store.py | 289 ++++++++++++++++++ .../proxy/proxy_server/test_proxy_config.py | 33 ++ 8 files changed, 848 insertions(+), 75 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 8c53216f77a..809a2086a06 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -41,6 +41,7 @@ litellm/proxy/_experimental/mcp_server/ sampling_handler.py # MCP sampling to LiteLLM completion flow elicitation_handler.py # MCP elicitation relay flow semantic_tool_filter.py # semantic filtering of available MCP tools + semantic_tool_index_store.py # Redis-shared embeddings + build lock for the semantic tool index tool_search.py # opt-in virtual tools (mcp_tool_search + mcp_tool_call) for large catalogs guardrail_translation/ handler.py # MCP guardrail result translation diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 2f88d8c361f..16398f4d035 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -8,17 +8,26 @@ import asyncio from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final +from mcp.types import Tool as MCPTool +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.exceptions import ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree +from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + CachingToolEncoder, + ToolVectorStore, + tool_vector_key, +) from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from litellm.router import Router + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder class SemanticToolFilterContextWindowError(Exception): @@ -58,6 +67,8 @@ class SemanticMCPToolFilter: top_k: int = 10, similarity_threshold: float = 0.3, enabled: bool = True, + vector_store: ToolVectorStore | None = None, + lock_poll_interval_s: float = 5.0, ): """ Initialize the semantic tool filter. @@ -68,6 +79,11 @@ class SemanticMCPToolFilter: top_k: Maximum number of tools to return similarity_threshold: Minimum similarity score for filtering enabled: Whether filtering is enabled + vector_store: Shared embedding store; when set the index builds + cross-pod in the background and requests pass through unfiltered + until startup_index_ready flips + lock_poll_interval_s: How often to re-check the store while another + pod holds the build lock """ self.enabled = enabled self.top_k = top_k @@ -78,10 +94,58 @@ class SemanticMCPToolFilter: self.context_window_error: str | None = None self._tool_map: dict[str, object] = {} # MCPTool objects or OpenAI function dicts self._index_sync_lock = asyncio.Lock() + self._vector_store = vector_store + self._lock_poll_interval_s = lock_poll_interval_s + self.embedding_identity = self._compute_embedding_identity() + self.startup_index_ready = vector_store is None + + def _compute_embedding_identity(self) -> str: + """embedding_model plus every deployment model it routes to, so repointing an alias invalidates the shared vectors.""" + try: + deployments: Final = self.router_instance.get_model_list(model_name=self.embedding_model) or () + deployment_maps: Final = TypeAdapter(tuple[Mapping[str, object], ...]).validate_python(deployments) + litellm_params: Final = tuple( + TypeAdapter(Mapping[str, object]).validate_python(deployment.get("litellm_params") or {}) + for deployment in deployment_maps + ) + deployment_models: Final = sorted( + {str(model) for model in (p.get("model") for p in litellm_params) if model} + ) + except (AttributeError, TypeError, KeyError, ValueError): + return self.embedding_model + if not deployment_models: + return self.embedding_model + return self.embedding_model + "|" + "|".join(deployment_models) + + def _new_encoder(self) -> "LiteLLMRouterEncoder": + """Encoder used for every index build; cached in Redis when a store is set.""" + from litellm.router_strategy.auto_router.litellm_encoder import ( + LiteLLMRouterEncoder, + ) + + if self._vector_store is not None: + return CachingToolEncoder( + litellm_router_instance=self.router_instance, + model_name=self.embedding_model, + score_threshold=self.similarity_threshold, + store=self._vector_store, + embedding_identity=self.embedding_identity, + ) + return LiteLLMRouterEncoder( + litellm_router_instance=self.router_instance, + model_name=self.embedding_model, + score_threshold=self.similarity_threshold, + ) @catalog_operation(global_manager) - async def build_router_from_mcp_registry(self) -> None: - """Build semantic router from all MCP tools in the registry (no auth checks).""" + async def build_router_from_mcp_registry(self, *, async_index: bool = False) -> None: + """Build semantic router from all MCP tools in the registry (no auth checks). + + With a vector store the shared Redis path runs: acquire the build lock + or wait for the holding pod's vectors, then embed only the misses. + Without one, the sync build blocks startup exactly as before; the + keyword arg lets a background task run the async (aadd) path instead. + """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -95,7 +159,7 @@ class SemanticMCPToolFilter: return # Fetch tools from all servers in parallel - all_tools: Final = [] + all_tools: Final[Sequence[object]] = [] for server_id, server in registry.items(): try: tools = await global_mcp_server_manager.get_tools_for_server(server_id) @@ -110,39 +174,66 @@ class SemanticMCPToolFilter: return verbose_logger.info("Fetched %s tools from %s MCP servers", len(all_tools), len(registry)) - self._build_router(all_tools) + if self._vector_store is not None or async_index: + await self._abuild_router(all_tools) + else: + self._build_router(all_tools) except Exception as e: verbose_logger.error("Failed to build router from MCP registry: %s", e) self.tool_router = None raise + finally: + self.startup_index_ready = True - def extract_tool_info(self, tool) -> tuple[str, str]: + def extract_tool_info(self, tool: object) -> tuple[str, str]: """Extract name and description from MCP tool or OpenAI function dict.""" name: str description: str if isinstance(tool, dict): # OpenAI function format - name = tool.get("name", "") - description = tool.get("description", name) + tool_dict: Final = TypeAdapter(Mapping[str, object]).validate_python(tool) + name = str(tool_dict.get("name", "")) + description = str(tool_dict.get("description", name)) else: # MCPTool object - name = str(tool.name) - description = str(tool.description) if tool.description else str(tool.name) + mcp_tool: Final = MCPTool.model_validate(tool, from_attributes=True) + name = str(mcp_tool.name) + description = str(mcp_tool.description) if mcp_tool.description else str(mcp_tool.name) return name, description _extract_tool_info = extract_tool_info - def _build_router(self, tools: list) -> None: - """Build semantic router with tools (MCPTool objects or OpenAI function dicts).""" - from semantic_router.routers import SemanticRouter + def _tools_to_routes_and_map(self, tools: Sequence[object]) -> tuple[Sequence[object], Mapping[str, object]]: from semantic_router.routers.base import Route - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) + extracted: Final = tuple((self.extract_tool_info(tool), tool) for tool in tools) + tool_map: Final = {name: tool for (name, _), tool in extracted} + routes: Final[Sequence[object]] = [ + Route( + name=name, + description=description, + utterances=[description], + score_threshold=self.similarity_threshold, + ) + for (name, description), _ in extracted + ] + return routes, tool_map + + def _handle_build_error(self, error: Exception) -> bool: + """True when the failure was recorded as a context-window overflow; False when the caller should re-raise.""" + verbose_logger.error("Failed to build semantic router: %s", error) + self.tool_router = None + if _is_context_window_error(error): + self.context_window_error = str(error) + return True + return False + + def _build_router(self, tools: Sequence[object]) -> None: + """Build semantic router with tools (MCPTool objects or OpenAI function dicts).""" + from semantic_router.routers import SemanticRouter if not tools: self.tool_router = None @@ -150,53 +241,95 @@ class SemanticMCPToolFilter: try: self.context_window_error = None - # Convert tools to routes - routes: Final = [] - self._tool_map = {} - - for tool in tools: - name, description = self.extract_tool_info(tool) - self._tool_map[name] = tool - - routes.append( - Route( - name=name, - description=description, - utterances=[description], - score_threshold=self.similarity_threshold, - ) - ) - + routes, tool_map = self._tools_to_routes_and_map(tools) self.tool_router = SemanticRouter( - routes=routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.router_instance, - model_name=self.embedding_model, - score_threshold=self.similarity_threshold, - ), + routes=list(routes), + encoder=self._new_encoder(), auto_sync="local", ) + self._tool_map = dict(tool_map) verbose_logger.info("Built semantic router with %s tools", len(routes)) except Exception as e: - verbose_logger.error("Failed to build semantic router: %s", e) - self.tool_router = None - if _is_context_window_error(e): - self.context_window_error = str(e) + if self._handle_build_error(e): return raise + async def _abuild_router(self, tools: Sequence[object]) -> None: + """ + Async index build for the background path. + + When a store is set, this first waits on the cross-pod build lock so + only one pod embeds; the caching encoder then fills any vectors the + holder did not write. + """ + from semantic_router.routers import SemanticRouter + + if not tools: + self.tool_router = None + return + + lock_held: Final = await self._await_store_or_lock(tools) if self._vector_store is not None else False + try: + self.context_window_error = None + routes, tool_map = self._tools_to_routes_and_map(tools) + encoder: Final = self._new_encoder() + router: Final = SemanticRouter( # pyright: ignore[reportUnknownVariableType] # semantic_router ships no type stubs + routes=[], + encoder=encoder, + init_async_index=True, + top_k=self.top_k, + ) + + # aadd's lazy init probes index dimensions through the sync encoder, + # so dimensions are seeded with an async probe to keep the build async. + dims_probe: Final = await encoder.aencode_queries(["test"]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + router.index.dimensions = len(dims_probe[0]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + await router.aadd(list(routes)) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + + self._tool_map = dict(tool_map) + self.tool_router = router + + verbose_logger.info("Built semantic router with %s tools via async index build", len(routes)) + + except Exception as e: + if self._handle_build_error(e): + return + raise + finally: + if lock_held and self._vector_store is not None: + await self._vector_store.release_build_lock() + + async def _await_store_or_lock(self, tools: Sequence[object]) -> bool: + """True when this pod holds the build lock after the wait; False when every vector arrived without it.""" + store: Final = self._vector_store + if store is None: + return False + + descriptions: Final = tuple(self.extract_tool_info(tool)[1] for tool in tools) + keys: Final = tuple(tool_vector_key(self.embedding_identity, text) for text in set(descriptions)) + + while {key for key in keys if key not in await store.get_many(keys)}: + if await store.try_acquire_build_lock(): + return True + await asyncio.sleep(self._lock_poll_interval_s) + return False + def _has_tools_missing_from_index(self, tools: Sequence[object]) -> bool: """Allocation-free check for any named tool not yet in the semantic index.""" - return any(name and name not in self._tool_map for name in (self.extract_tool_info(t)[0] for t in tools)) + return any( + name and (name not in self._tool_map or self.extract_tool_info(self._tool_map[name])[1] != description) + for name, description in (self.extract_tool_info(t) for t in tools) + ) def _tools_missing_from_index(self, tools: Sequence[object]) -> Mapping[str, object]: """Map name -> tool for every named tool not yet in the semantic index.""" + infos: Final = ((self.extract_tool_info(t), t) for t in tools) return { name: tool - for name, tool in ((self.extract_tool_info(t)[0], t) for t in tools) - if name and name not in self._tool_map + for (name, description), tool in infos + if name and (name not in self._tool_map or self.extract_tool_info(self._tool_map[name])[1] != description) } async def _ensure_tools_indexed(self, available_tools: Sequence[object]) -> None: @@ -216,11 +349,6 @@ class SemanticMCPToolFilter: description cannot poison the filter for other users on the worker. """ from semantic_router.routers import SemanticRouter - from semantic_router.routers.base import Route - - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) if not self._has_tools_missing_from_index(available_tools): return @@ -230,32 +358,21 @@ class SemanticMCPToolFilter: if not missing: return - descriptions: Final = {name: self.extract_tool_info(tool)[1] for name, tool in missing.items()} - routes: Final = [ - Route( - name=name, - description=description, - utterances=[description], - score_threshold=self.similarity_threshold, - ) - for name, description in descriptions.items() - ] + routes, _ = self._tools_to_routes_and_map(tuple(missing.values())) if self.tool_router is None: router: Final = SemanticRouter( routes=[], - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.router_instance, - model_name=self.embedding_model, - score_threshold=self.similarity_threshold, - ), + encoder=self._new_encoder(), auto_sync="local", top_k=self.top_k, ) - await router.aadd(routes) + await router.aadd(list(routes)) self.tool_router = router else: - await self.tool_router.aadd(routes) + for name in (name for name in missing if name in self._tool_map): + await self.tool_router.adelete(name) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + await self.tool_router.aadd(list(routes)) self._tool_map.update(missing) verbose_logger.info( diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py new file mode 100644 index 00000000000..9429ecb2989 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -0,0 +1,172 @@ +""" +Redis-shared vector store for the MCP semantic tool index. + +Lets every pod reuse embeddings another pod already computed instead of +re-embedding every tool description at startup. +""" + +import base64 +import sys +from array import array +from collections.abc import Mapping, Sequence +from hashlib import sha256 +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter + +from litellm._logging import verbose_logger +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.router import Router + +_KEY_PREFIX: Final = "litellm:mcp_semantic_tool_index:v1:" +_BUILD_LOCK_ID: Final = "mcp_semantic_tool_index_build" +_BUILD_LOCK_TTL_S: Final = 300 +_VECTOR_TTL_S: Final = 30 * 24 * 60 * 60 +_PUT_CHUNK_SIZE: Final = 50 + + +def tool_vector_key(embedding_identity: str, text: str) -> str: + """Cache key for one embedded description; identity changes invalidate every entry.""" + return _KEY_PREFIX + sha256((embedding_identity + "\0" + text).encode("utf-8")).hexdigest() + + +class ToolVectorStore(Protocol): + """Shared embedding cache plus a cross-pod build lock.""" + + async def get_many(self, keys: Sequence[str]) -> Mapping[str, tuple[float, ...]]: + """Return only the keys that hit.""" + ... + + async def put_many(self, vectors: Mapping[str, Sequence[float]]) -> None: ... + + async def try_acquire_build_lock(self) -> bool: ... + + async def release_build_lock(self) -> None: ... + + +class RedisToolVectorStore: + """ToolVectorStore over RedisCache; every Redis failure degrades to a local embed.""" + + def __init__(self, redis_cache: "RedisCache", lock_manager: PodLockManager) -> None: + self._redis_cache = redis_cache + self._lock_manager = lock_manager + self._degraded = False + + async def get_many(self, keys: Sequence[str]) -> Mapping[str, tuple[float, ...]]: + try: + raw: Final = TypeAdapter(Mapping[str, object]).validate_python( + await self._redis_cache.async_batch_get_cache(list(keys)) # pyright: ignore[reportUnknownMemberType] # RedisCache declares a partially untyped signature + ) + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + self._degraded = True + verbose_logger.warning("MCP semantic tool index store get_many failed; embedding locally: %s", e) + return {} + return {key: vector for key in keys if (vector := _decode_vector(raw.get(key))) is not None} + + async def put_many(self, vectors: Mapping[str, Sequence[float]]) -> None: + try: + entries: Final = tuple((key, _encode_vector(vector)) for key, vector in vectors.items()) + for start in range(0, len(entries), _PUT_CHUNK_SIZE): + await self._redis_cache.async_set_cache_pipeline( # pyright: ignore[reportUnknownMemberType] # RedisCache declares a partially untyped signature + entries[start : start + _PUT_CHUNK_SIZE], ttl=_VECTOR_TTL_S + ) + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + self._degraded = True + verbose_logger.warning("MCP semantic tool index store put_many failed; vectors stay local: %s", e) + + async def try_acquire_build_lock(self) -> bool: + # PodLockManager returns False both when another pod holds the lock and + # when Redis errors internally, so a degraded store cannot tell them + # apart; err toward embedding locally rather than waiting on writes + # that a dead Redis can never deliver. + try: + acquired: Final = await self._lock_manager.acquire_lock( + _BUILD_LOCK_ID, ttl=_BUILD_LOCK_TTL_S, allow_reentrant=True + ) + return bool(acquired) or self._degraded + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + self._degraded = True + verbose_logger.warning("MCP semantic tool index store lock acquire failed; embedding locally: %s", e) + return True + + async def release_build_lock(self) -> None: + try: + await self._lock_manager.release_lock(_BUILD_LOCK_ID) + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + verbose_logger.warning("MCP semantic tool index store lock release failed: %s", e) + + +def _encode_vector(vector: Sequence[float]) -> str: + packed: Final = array("f", vector) + if sys.byteorder == "big": + packed.byteswap() + return base64.b64encode(packed.tobytes()).decode("ascii") + + +def _decode_vector(value: object) -> tuple[float, ...] | None: + try: + if not isinstance(value, (str, bytes, bytearray)): + return None + raw: Final = base64.b64decode(value) + vector: Final = array("f") + vector.frombytes(raw) + if sys.byteorder == "big": + vector.byteswap() + return tuple(vector) + except Exception: # noqa: BLE001 # corrupt payloads are treated as a plain cache miss + return None + + +class CachingToolEncoder(LiteLLMRouterEncoder): + """LiteLLMRouterEncoder that serves document embeddings from a ToolVectorStore. + + Queries are never cached; only add/aadd's document utterances go through + the store. Duplicate texts share one key and one embed. + """ + + _tool_vector_store: ToolVectorStore + _embedding_identity: str + + def __init__( + self, + *, + store: ToolVectorStore, + embedding_identity: str, + litellm_router_instance: "Router", + model_name: str, + score_threshold: float | None = None, + ) -> None: + super().__init__( + litellm_router_instance=litellm_router_instance, + model_name=model_name, + score_threshold=score_threshold, + ) + self._tool_vector_store = store + self._embedding_identity = embedding_identity + + def encode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls encode_documents without kwargs + self, docs: Sequence[str] + ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract + return super().encode_documents(list(docs)) + + async def aencode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls aencode_documents without kwargs + self, docs: Sequence[str] + ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract + unique_docs: Final = tuple(dict.fromkeys(docs)) + keys: Final = tuple(tool_vector_key(self._embedding_identity, text) for text in unique_docs) + hits: Final = await self._tool_vector_store.get_many(keys) + + miss_pairs: Final = tuple((text, key) for text, key in zip(unique_docs, keys) if key not in hits) + embedded: Final = dict(hits) + if miss_pairs: + vectors: Final = await super().aencode_documents([text for text, _ in miss_pairs]) + new_vectors: Final = {key: tuple(vector) for (_, key), vector in zip(miss_pairs, vectors)} + await self._tool_vector_store.put_many(new_vectors) + embedded.update(new_vectors) + + vectors_by_text: Final = dict(zip(unique_docs, (embedded[key] for key in keys))) + return [list(vectors_by_text[text]) for text in docs] diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index d33811d3dea..8cedea69c6f 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,6 +5,8 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ +import asyncio +import importlib.util from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional @@ -24,13 +26,39 @@ from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( if TYPE_CHECKING: from litellm.caching.caching import DualCache + from litellm.caching.redis_cache import RedisCache from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( SemanticMCPToolFilter, ) + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ToolVectorStore from litellm.proxy._types import UserAPIKeyAuth from litellm.router import Router +def _make_vector_store(redis_cache: Optional["RedisCache"]) -> Optional["ToolVectorStore"]: + """Redis-backed shared vector store, or None when no Redis is configured.""" + if redis_cache is None: + return None + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + ) + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import ( + PodLockManager, + ) + + return RedisToolVectorStore(redis_cache, PodLockManager(redis_cache)) + + +async def _run_deferred_index_build(semantic_filter: "SemanticMCPToolFilter") -> None: + """Background task for a shared (Redis-backed) index build; swallows failures so requests keep passing through.""" + try: + await semantic_filter.build_router_from_mcp_registry(async_index=True) + except Exception: # noqa: BLE001 # background task must not die; requests keep passing through + verbose_proxy_logger.exception( + "MCP semantic tool index deferred build failed; request-time indexing remains active" + ) + + class SemanticToolFilterConfig(TypedDict, total=False): enabled: ReadOnly[bool] embedding_model: ReadOnly[str] @@ -69,6 +97,7 @@ class SemanticToolFilterHook(CustomLogger): """ super().__init__() self.filter = semantic_filter + self.index_build_task: asyncio.Task[None] | None = None verbose_proxy_logger.debug( "Initialized SemanticToolFilterHook with filter: enabled=%s, top_k=%s", @@ -314,8 +343,11 @@ class SemanticToolFilterHook(CustomLogger): if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug("Detected litellm_proxy MCP references, expanding before semantic filtering") - if not self.filter.enabled: - verbose_proxy_logger.debug("Semantic filter disabled, leaving MCP references untouched") + if not self.filter.enabled or not self.filter.startup_index_ready: + verbose_proxy_logger.debug( + "Semantic filter %s, leaving MCP references untouched", + "disabled" if not self.filter.enabled else "index still building", + ) return None try: @@ -360,8 +392,11 @@ class SemanticToolFilterHook(CustomLogger): verbose_proxy_logger.debug("No messages in request, skipping semantic filter") return None - if not self.filter.enabled: - verbose_proxy_logger.debug("Semantic filter disabled, skipping") + if not self.filter.enabled or not self.filter.startup_index_ready: + verbose_proxy_logger.debug( + "Semantic filter %s, skipping", + "disabled" if not self.filter.enabled else "index still building", + ) return None try: @@ -476,6 +511,7 @@ class SemanticToolFilterHook(CustomLogger): async def initialize_from_config( config: SemanticToolFilterConfig | None, llm_router: Optional["Router"], + redis_cache: Optional["RedisCache"] = None, ) -> Optional["SemanticToolFilterHook"]: """ Initialize semantic tool filter from proxy config. @@ -483,6 +519,9 @@ class SemanticToolFilterHook(CustomLogger): Args: config: Proxy configuration dict (litellm_settings.mcp_semantic_tool_filter) llm_router: LiteLLM router instance for embeddings + redis_cache: Usage-tracking Redis cache; when configured the tool + index is built in the background with embeddings shared + cross-pod through it Returns: SemanticToolFilterHook instance if enabled, None otherwise @@ -500,6 +539,9 @@ class SemanticToolFilterHook(CustomLogger): return None try: + if redis_cache is not None and importlib.util.find_spec("semantic_router") is None: + raise ImportError("No module named 'semantic_router'") + embedding_model: Final = config.get("embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL) top_k: Final = config.get("top_k", DEFAULT_MCP_SEMANTIC_FILTER_TOP_K) similarity_threshold = config.get("similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD) @@ -510,18 +552,24 @@ class SemanticToolFilterHook(CustomLogger): top_k=top_k, similarity_threshold=similarity_threshold, enabled=True, + vector_store=_make_vector_store(redis_cache), ) - # Build router from MCP registry on startup - await semantic_filter.build_router_from_mcp_registry() - hook: Final = SemanticToolFilterHook(semantic_filter) + if redis_cache is not None: + # Index builds in the background; requests pass through unfiltered until it finishes. + hook.index_build_task = asyncio.create_task(_run_deferred_index_build(semantic_filter)) + else: + # Build router from MCP registry on startup + await semantic_filter.build_router_from_mcp_registry() + verbose_proxy_logger.info( - "✅ MCP Semantic Tool Filter enabled: embedding_model=%s, top_k=%s, similarity_threshold=%s", + "✅ MCP Semantic Tool Filter enabled: embedding_model=%s, top_k=%s, similarity_threshold=%s, shared_index=%s", embedding_model, top_k, similarity_threshold, + redis_cache is not None, ) return hook diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dac46855fd6..7d04b4f016f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8560,6 +8560,11 @@ class ProxyConfig: if active_hooks: for active_hook in active_hooks: if isinstance(active_hook, SemanticToolFilterHook): + if active_hook.index_build_task is not None: + verbose_proxy_logger.debug( + "Semantic filter settings unchanged, index build already started; skipping reinitialization" + ) + return if active_hook.filter is not None and active_hook.filter.tool_router is not None: verbose_proxy_logger.debug( "Semantic filter settings unchanged, skipping reinitialization" @@ -8578,6 +8583,7 @@ class ProxyConfig: hook: Final = await SemanticToolFilterHook.initialize_from_config( config=mcp_semantic_filter_config, llm_router=llm_router, + redis_cache=redis_usage_cache, ) if hook: litellm.logging_callback_manager.add_litellm_callback(hook) @@ -10596,6 +10602,7 @@ class ProxyStartupEvent: hook: Final = await SemanticToolFilterHook.initialize_from_config( config=mcp_semantic_filter_config, llm_router=llm_router, + redis_cache=redis_usage_cache, ) if hook: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index f0b4e94f72f..1ba1b33b094 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2239,3 +2239,109 @@ async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves() ) assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_semantic_filter_reindexes_tool_when_description_changes(): + """ + Regression test: a tool that keeps its name but gets a new description + must be re-embedded, otherwise the index keeps serving the stale vector. + """ + recorded_inputs = [] + filter_instance = _make_keyword_filter(recorded_inputs) + original = MCPTool( + name="srv-a-get_thing", + description="get the weather for a city", + inputSchema={"type": "object"}, + ) + edited = MCPTool( + name="srv-a-get_thing", + description="get a Linear ticket by id", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + filtered = await filter_instance.filter_tools( + query="linear ticket", + available_tools=[edited], + ) + + assert [edited] == filtered + assert any("get a Linear ticket by id" in text for batch in recorded_inputs for text in batch) + indexed_tool = filter_instance._tool_map["srv-a-get_thing"] + assert filter_instance.extract_tool_info(indexed_tool)[1] == "get a Linear ticket by id" + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_hook_returns_before_shared_index_build_finishes(): + """ + With Redis configured the hook registers immediately and the shared + index build runs as a background task that finishes once the store + responds; without Redis the build stays inline. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + gate = asyncio.Event() + + class _GatedRedis: + def __init__(self): + self.backing = {} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + await gate.wait() + return {key: self.backing[key] for key in key_list if key in self.backing} + + async def async_set_cache(self, key, value, **kwargs): + self.backing[key] = value + return True + + async def async_get_cache(self, key): + return self.backing.get(key) + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + self.backing.update(dict(cache_list)) + + recorded_inputs = [] + tool = MCPTool( + name="srv-a-get_thing", + description="get a Linear ticket by id", + inputSchema={"type": "object"}, + ) + with ( + patch.object(global_mcp_server_manager, "get_registry", return_value={"srv-a": object()}), + patch.object( + global_mcp_server_manager, "get_tools_for_server", new=AsyncMock(return_value=[tool]) + ), + ): + hook = await SemanticToolFilterHook.initialize_from_config( + config={"enabled": True}, + llm_router=_make_keyword_embedding_router(recorded_inputs), + redis_cache=_GatedRedis(), + ) + + assert hook is not None + assert hook.index_build_task is not None + await asyncio.sleep(0) + assert not hook.index_build_task.done() + assert not hook.filter.startup_index_ready + + gate.set() + await asyncio.wait_for(hook.index_build_task, timeout=30) + assert hook.filter.startup_index_ready + assert hook.filter.tool_router is not None + + inline_hook = await SemanticToolFilterHook.initialize_from_config( + config={"enabled": True}, + llm_router=_make_keyword_embedding_router([]), + redis_cache=None, + ) + assert inline_hook is not None + assert inline_hook.index_build_task is None + assert inline_hook.filter.startup_index_ready + assert inline_hook.filter.tool_router is not None diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py new file mode 100644 index 00000000000..0530de440ad --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py @@ -0,0 +1,289 @@ +""" +Unit tests for the Redis-shared MCP semantic tool index store. + +Covers the float32 base64 round trip, degradation on Redis failure, the +cross-pod build lock, and the CachingToolEncoder serving document +embeddings from the store instead of re-embedding. +""" + +import sys +from unittest.mock import Mock + +import pytest + +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + + +from mcp.types import Tool as MCPTool + + +requires_semantic_router = pytest.mark.skipif( + sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" +) + + +def _fake_embedding_router(recorded_inputs): + """Router stand-in returning deterministic 2-dim keyword one-hots, like the filter tests.""" + from litellm.types.utils import Embedding, EmbeddingResponse + + def _vector(text): + return [1.0, 0.0] if "ticket" in text.lower() else [0.0, 1.0] + + def mock_embedding_sync(*args, **kwargs): + texts = kwargs["input"] + recorded_inputs.append(list(texts)) + return EmbeddingResponse( + data=[Embedding(embedding=_vector(t), index=i, object="embedding") for i, t in enumerate(texts)], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync(*args, **kwargs) + + router = Mock() + router.embedding = mock_embedding_sync + router.aembedding = mock_embedding_async + return router + + +class _InMemoryToolVectorStore: + def __init__(self, vectors=None): + self.vectors = dict(vectors or {}) + self.get_many_calls = 0 + self.acquire_calls = 0 + self.release_calls = 0 + self.lock_result = True + self._fill_on_call = None + + def fill_on_get_many_call(self, call_number): + self._fill_on_call = call_number + + async def get_many(self, keys): + self.get_many_calls += 1 + if self._fill_on_call is not None and self.get_many_calls >= self._fill_on_call: + self.vectors.update(self._pending_fill) + return {key: self.vectors[key] for key in keys if key in self.vectors} + + def seed_from(self, keys_vectors): + self._pending_fill = dict(keys_vectors) + + async def put_many(self, vectors): + self.vectors.update({k: tuple(v) for k, v in vectors.items()}) + + async def try_acquire_build_lock(self): + self.acquire_calls += 1 + return self.lock_result + + async def release_build_lock(self): + self.release_calls += 1 + + +class _RaisingRedisCache: + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + raise ConnectionError("redis is down") + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + raise ConnectionError("redis is down") + + async def async_get_cache(self, key): + raise ConnectionError("redis is down") + + async def async_set_cache(self, key, value, **kwargs): + raise ConnectionError("redis is down") + + +class _DictRedisCache: + def __init__(self): + self.backing = {} + self.pipeline_calls = 0 + + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + return {key: self.backing[key] for key in key_list if key in self.backing} + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + self.pipeline_calls += 1 + self.backing.update(dict(cache_list)) + + +def _tool(name, description): + return MCPTool(name=name, description=description, inputSchema={"type": "object"}) + + +def _make_filter(recorded_inputs, store=None, embedding_model="text-embedding-3-small"): + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + + return SemanticMCPToolFilter( + embedding_model=embedding_model, + litellm_router_instance=_fake_embedding_router(recorded_inputs), + top_k=3, + similarity_threshold=0.3, + enabled=True, + vector_store=store, + lock_poll_interval_s=0.0, + ) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_store_put_get_round_trip_within_float32_tolerance(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + tool_vector_key, + ) + + redis = _DictRedisCache() + store = RedisToolVectorStore(redis, Mock()) + + key = tool_vector_key("embed-a", "get a ticket") + await store.put_many({key: [0.6000000238, -0.8123456789, 3.14]}) + + hits = await store.get_many([key]) + assert key in hits + assert list(hits[key]) == pytest.approx([0.6000000238, -0.8123456789, 3.14], abs=1e-6) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_store_treats_corrupt_value_as_miss(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + tool_vector_key, + ) + + redis = _DictRedisCache() + key = tool_vector_key("embed-a", "get a ticket") + redis.backing[key] = "!!! not base64 / not float32 bytes !!!" + store = RedisToolVectorStore(redis, Mock()) + + assert await store.get_many([key]) == {} + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_raising_redis_still_builds_the_index(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + ) + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + + redis = _RaisingRedisCache() + recorded = [] + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = RedisToolVectorStore(redis, PodLockManager(redis)) + filter_instance = _make_filter(recorded, store=store) + + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + assert [batch for batch in recorded if batch != ["test"]] == [ + ["get a ticket", "weather forecast"] + ] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_build_with_warm_store_never_embeds_descriptions(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + identity = "text-embedding-3-small" + warm = _InMemoryToolVectorStore( + {tool_vector_key(identity, description): [0.0, 1.0] for _, description in ((t.name, t.description) for t in tools)} + ) + recorded = [] + filter_instance = _make_filter(recorded, store=warm) + + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + # Only the dims probe embeds; the 300-tool warm path makes zero embedding calls. + assert recorded == [["test"]] + assert warm.acquire_calls == 0 + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_cold_store_embeds_once_then_second_filter_embeds_nothing(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + + first_recorded = [] + first = _make_filter(first_recorded, store=store) + await first._abuild_router(tools) + assert first.tool_router is not None + assert [batch for batch in first_recorded if batch != ["test"]] == [ + ["get a ticket", "weather forecast"] + ] + + second_recorded = [] + second = _make_filter(second_recorded, store=store) + await second._abuild_router(tools) + assert second.tool_router is not None + assert second_recorded == [["test"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_edited_description_reembeds_only_that_tool(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + recorded = [] + filter_instance = _make_filter(recorded, store=store) + await filter_instance._abuild_router(tools) + + recorded.clear() + edited = [_tool("srv-a_get_ticket", "close a support ticket fast"), _tool("srv-b_get_weather", "weather forecast")] + await filter_instance._abuild_router(edited) + + assert [batch for batch in recorded if batch != ["test"]] == [["close a support ticket fast"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_different_embedding_identity_reembeds_everything(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + + first_recorded = [] + await _make_filter(first_recorded, store=store, embedding_model="text-embedding-3-small")._abuild_router(tools) + + second_recorded = [] + second = _make_filter(second_recorded, store=store, embedding_model="text-embedding-3-large") + await second._abuild_router(tools) + + assert second.tool_router is not None + assert [batch for batch in second_recorded if batch != ["test"]] == [ + ["get a ticket", "weather forecast"] + ] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_held_by_other_pod_waits_for_shared_vectors(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + store.lock_result = False + identity = "text-embedding-3-small" + # The "holder" pod's vectors appear on the third get_many poll. + store.seed_from({tool_vector_key(identity, t.description): [0.0, 1.0] for t in tools}) + store.fill_on_get_many_call(3) + + recorded = [] + filter_instance = _make_filter(recorded, store=store) + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + assert store.acquire_calls >= 1 + assert store.get_many_calls >= 3 + assert recorded == [["test"]] diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 720695b0f42..66b2f2787ac 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -5607,3 +5607,36 @@ async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, return with pytest.raises(ValidationError): await ProxyConfig().load_config(router=None, config_file_path=str(config)) + + +@pytest.mark.asyncio +async def test_proxy_config_semantic_filter_db_sync_skips_reinit_while_index_build_in_flight(monkeypatch): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + settings: Final = { + "enabled": True, + "embedding_model": "embed", + "top_k": 3, + "similarity_threshold": 0.5, + } + gate: Final = asyncio.Event() + build_task: Final = asyncio.create_task(gate.wait()) + hook: Final = MagicMock(spec=SemanticToolFilterHook) + hook.index_build_task = build_task + hook.filter = SimpleNamespace(tool_router=None) + removed: Final = MagicMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_config_param", + AsyncMock(return_value=SimpleNamespace(param_value={"mcp_semantic_tool_filter": settings})), + ) + monkeypatch.setattr(litellm.logging_callback_manager, "get_custom_loggers_for_type", lambda _: [hook]) + monkeypatch.setattr(litellm.logging_callback_manager, "remove_callbacks_by_type", removed) + pc: Final = ProxyConfig() + pc._last_semantic_filter_config = settings + try: + await pc._init_semantic_filter_settings_in_db(prisma_client=MagicMock()) + + removed.assert_not_called() + finally: + gate.set() + await build_task