diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 8c53216f77a..2abab418446 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -41,6 +41,8 @@ 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 + semantic_tool_encoder.py # CachingToolEncoder bridging the store and LiteLLMRouterEncoder 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_encoder.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_encoder.py new file mode 100644 index 00000000000..338b5f06459 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_encoder.py @@ -0,0 +1,83 @@ +""" +Caching encoder for the shared MCP semantic tool index. + +Lives apart from semantic_tool_index_store because LiteLLMRouterEncoder pulls +in semantic_router, an optional extra that must stay out of the store's import +chain. +""" + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + ToolVectorStore, + tool_vector_key, +) +from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + +if TYPE_CHECKING: + from litellm.router import Router + + +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. Cached vectors + whose length disagrees with a live probe are treated as misses, so a + stale entry can never feed the index a wrong-shape vector. + """ + + _tool_vector_store: ToolVectorStore + _embedding_identity: str + _dimensions: int | None + + 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 + self._dimensions = None + + async def adimensions(self) -> int: + if self._dimensions is None: + probe: Final = await self.aencode_queries(["test"]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + self._dimensions = len(probe[0]) + return self._dimensions + + 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[Mapping[str, tuple[float, ...]]] = await self._tool_vector_store.get_many(keys) + valid_hits: Final = ( + {key: vector for key, vector in hits.items() if len(vector) == await self.adimensions()} if hits else {} + ) + + miss_pairs: Final = tuple((text, key) for text, key in zip(unique_docs, keys) if key not in valid_hits) + embedded: Final = dict(valid_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/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 2f88d8c361f..6b2b93f2626 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -8,17 +8,25 @@ 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 ( + 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): @@ -48,6 +56,13 @@ def _is_context_window_error(error: BaseException | None) -> bool: ) +def _embedding_params_suffix(litellm_params: Mapping[str, object]) -> str: + """Deployment traits that make cached vectors non-interchangeable across endpoints.""" + return (f"@dimensions={litellm_params['dimensions']}" if litellm_params.get("dimensions") is not None else "") + ( + f"@api_base={litellm_params['api_base']}" if litellm_params.get("api_base") is not None else "" + ) + + class SemanticMCPToolFilter: """Filters MCP tools using semantic similarity to reduce context window size.""" @@ -58,6 +73,9 @@ 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, + lock_wait_timeout_s: float = 300.0, ): """ Initialize the semantic tool filter. @@ -68,6 +86,13 @@ 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 + lock_wait_timeout_s: Give up waiting for another pod's vectors after + this long and embed locally (matches the build-lock TTL) """ self.enabled = enabled self.top_k = top_k @@ -78,10 +103,63 @@ 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._lock_wait_timeout_s = lock_wait_timeout_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) + _embedding_params_suffix(params) + for model, params in ((p.get("model"), p) 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.proxy._experimental.mcp_server.semantic_tool_encoder import CachingToolEncoder + 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).""" + """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. + """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -95,7 +173,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 +188,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: + 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,112 +255,161 @@ 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 - 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)) - - 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.""" - 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 - } - - async def _ensure_tools_indexed(self, available_tools: Sequence[object]) -> None: + async def _abuild_router(self, tools: Sequence[object]) -> None: """ - Index request-time tools the startup build never saw. + Async index build for the background path. - The startup index lists every registered MCP server WITHOUT per-user - credentials, so servers requiring per-user auth (interactive OAuth - tokens, user-scoped env vars) contribute zero routes. Tools reaching - the filter came through an authenticated expansion; without indexing - them here they can never be selected, so requests either bypass - filtering entirely (N->N) or lose every tool to unrelated matches. - - Runs async-only (no synchronous embedding on the request path) and - never writes shared error state: an embedding failure here raises and - is scoped to the requesting call, so one request's oversized tool - description cannot poison the filter for other users on the worker. + 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 - from semantic_router.routers.base import Route - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, + from litellm.proxy._experimental.mcp_server.semantic_tool_encoder import CachingToolEncoder + + 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. + dimensions: Final = ( + await encoder.adimensions() + if isinstance(encoder, CachingToolEncoder) + else len((await encoder.aencode_queries(["test"]))[0]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + ) + router.index.dimensions = dimensions # 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_set: Final = frozenset(tool_vector_key(self.embedding_identity, text) for text in set(descriptions)) + + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + self._lock_wait_timeout_s + while keys_set - (await store.get_many(tuple(keys_set))).keys(): + if await store.try_acquire_build_lock(): + return True + if loop.time() >= deadline: + verbose_logger.warning( + "MCP semantic tool index wait exceeded %.1fs; embedding locally", + self._lock_wait_timeout_s, + ) + return False + await asyncio.sleep(self._lock_poll_interval_s) + return False + + def _has_tools_missing_from_index(self, tools: Sequence[object], *, reindex_changed_descriptions: bool) -> bool: + """Allocation-free check for any named tool not yet in the semantic index.""" + return any( + name + and ( + name not in self._tool_map + or (reindex_changed_descriptions and self.extract_tool_info(self._tool_map[name])[1] != description) + ) + for name, description in (self.extract_tool_info(t) for t in tools) ) - if not self._has_tools_missing_from_index(available_tools): + def _tools_missing_from_index( + self, tools: Sequence[object], *, reindex_changed_descriptions: bool + ) -> 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, description), tool in infos + if name + and ( + name not in self._tool_map + or (reindex_changed_descriptions and self.extract_tool_info(self._tool_map[name])[1] != description) + ) + } + + async def _ensure_tools_indexed( + self, available_tools: Sequence[object], *, reindex_changed_descriptions: bool + ) -> None: + """ + Index request-time tools missing from the startup index. + + With ``reindex_changed_descriptions`` a tool whose description changed + also counts as missing: its old route is dropped before re-embedding so + a failed embed leaves it "missing" and the next request retries. + """ + from semantic_router.routers import SemanticRouter + + if not self._has_tools_missing_from_index( + available_tools, reindex_changed_descriptions=reindex_changed_descriptions + ): return async with self._index_sync_lock: - missing: Final = self._tools_missing_from_index(available_tools) + missing: Final = self._tools_missing_from_index( + available_tools, reindex_changed_descriptions=reindex_changed_descriptions + ) 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 + self._tool_map.pop(name) + await self.tool_router.aadd(list(routes)) self._tool_map.update(missing) verbose_logger.info( @@ -267,6 +421,8 @@ class SemanticMCPToolFilter: query: str, available_tools: list[Any], top_k: int | None = None, + *, + reindex_changed_descriptions: bool = False, ) -> list[Any]: """ Filter tools semantically based on query. @@ -298,7 +454,7 @@ class SemanticMCPToolFilter: # Run semantic filtering try: - await self._ensure_tools_indexed(available_tools) + await self._ensure_tools_indexed(available_tools, reindex_changed_descriptions=reindex_changed_descriptions) if self.tool_router is None: verbose_logger.warning("Semantic router could not be built from the request's tools") 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..212bee108f8 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -0,0 +1,134 @@ +""" +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._internal_context import with_service_target +from litellm._logging import verbose_logger +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + +_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 +_CACHE_TARGET: Final = "mcp_semantic_tool_index" + + +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 + + @with_service_target(_CACHE_TARGET) + 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} + + @with_service_target(_CACHE_TARGET) + 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) + + @with_service_target(_CACHE_TARGET) + async def try_acquire_build_lock(self) -> bool: + # PodLockManager returns False on Redis errors too; once degraded, embed locally instead of waiting. + try: + acquired: Final = await self._lock_manager.acquire_lock( + _BUILD_LOCK_ID, ttl=_BUILD_LOCK_TTL_S, allow_reentrant=True + ) + held_by: Final = ( + None + if acquired or self._degraded + else await self._redis_cache.async_get_cache( + PodLockManager.get_redis_lock_key(_BUILD_LOCK_ID) # pyright: ignore[reportUnknownArgumentType] # 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 lock acquire failed; embedding locally: %s", e) + return True + if acquired or self._degraded: + return True + if held_by is not None: + return False + self._degraded = True + verbose_logger.warning("MCP semantic tool index store lock write failed with no holder; embedding locally") + 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 diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index d33811d3dea..3682213309b 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() + 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", @@ -165,7 +194,9 @@ class SemanticToolFilterHook(CustomLogger): verbose_proxy_logger.debug("No user query found, skipping semantic filter on expanded MCP tools") return expanded_tools - return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools) + return await self.filter.filter_tools( + query=user_query, available_tools=expanded_tools, reindex_changed_descriptions=True + ) def _selected_tool_names(self, filtered_tools: Sequence[object]) -> list[str]: """Names of the semantically selected tools, as produced by the MCP expansion.""" @@ -314,8 +345,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 +394,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 +513,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 +521,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 +541,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 +554,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 1c7cf54ab4d..eab6c872064 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8533,6 +8533,41 @@ class ProxyConfig: db_values: Final = self._prepared_db_settings_values("litellm_settings", raw_settings) self._apply_litellm_settings_db_values(db_values) + def _semantic_filter_settings_unchanged(self, mcp_semantic_filter_config: object, hook_type: type) -> bool: + """True when the stored settings match the in-memory ones and a live hook already covers them.""" + import litellm + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + if not hasattr(self, "_last_semantic_filter_config"): + return False + if self._last_semantic_filter_config != mcp_semantic_filter_config: + return False + for active_hook in litellm.logging_callback_manager.get_custom_loggers_for_type(hook_type): + if not isinstance(active_hook, SemanticToolFilterHook): + continue + if active_hook.index_build_task is not None and not active_hook.index_build_task.done(): + verbose_proxy_logger.debug( + "Semantic filter settings unchanged, index build already started; skipping reinitialization" + ) + return True + 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") + return True + verbose_proxy_logger.info( + "Semantic filter settings unchanged, but hook is missing or uninitialized. Reinitializing." + ) + return False + + @staticmethod + def _cancel_orphaned_semantic_filter_index_builds() -> None: + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + for active_hook in litellm.logging_callback_manager.get_custom_loggers_for_type(SemanticToolFilterHook): + if not isinstance(active_hook, SemanticToolFilterHook): + continue + if (build_task := active_hook.index_build_task) is not None and not build_task.done(): + build_task.cancel() + async def init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient) -> None: """ Initialize MCP semantic filter settings from database. @@ -8559,21 +8594,10 @@ class ProxyConfig: return # Check if settings have changed (compare with in-memory state) - if hasattr(self, "_last_semantic_filter_config"): - if self._last_semantic_filter_config == mcp_semantic_filter_config: - # If hook is missing or router isn't built yet, reinitialize anyway - active_hooks = litellm.logging_callback_manager.get_custom_loggers_for_type(SemanticToolFilterHook) - if active_hooks: - for active_hook in active_hooks: - if isinstance(active_hook, SemanticToolFilterHook): - 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" - ) - return - verbose_proxy_logger.info( - "Semantic filter settings unchanged, but hook is missing or uninitialized. Reinitializing." - ) + if self._semantic_filter_settings_unchanged(mcp_semantic_filter_config, SemanticToolFilterHook): + return + + self._cancel_orphaned_semantic_filter_index_builds() # Remove old hooks using logging callback manager litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, SemanticToolFilterHook) @@ -8584,6 +8608,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) @@ -10602,6 +10627,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/integration/mcp/test_mcp_semantic_filter_startup.py b/tests/integration/mcp/test_mcp_semantic_filter_startup.py new file mode 100644 index 00000000000..64ed44c6a90 --- /dev/null +++ b/tests/integration/mcp/test_mcp_semantic_filter_startup.py @@ -0,0 +1,513 @@ +import json +import os +import sys +import threading +import time +import uuid +from collections import Counter +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.client import GATEWAY_LIMITS, Gateway, eventually +from integration._support.mcp import ScriptedTool, scripted_peer, text_result +from integration._support.process import ( + DB_PUSH, + OwnedProxy, + _launch, + _proxy_environment, + _proxy_root, + _stop_launch, + owned_proxy_process, +) +from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.wire import Reply, Request, wire_server + +PROMPT: Final = "please refund the invoice for ACME" +MCP_REF: Final = { + "type": "mcp", + "server_url": "litellm_proxy", + "server_label": "litellm", + "require_approval": "never", +} +SEMANTIC_SETTINGS: Final = { + "enabled": True, + "embedding_model": "semantic-embed", + "top_k": 3, + "similarity_threshold": 0.5, +} +NO_REDIS_ENV: Final = ( + "REDIS_HOST", + "REDIS_PORT", + "REDIS_URL", + "REDIS_USERNAME", + "REDIS_PASSWORD", + "REDIS_CLUSTER_NODES", + "REDIS_SENTINEL_NODES", + "REDIS_SSL", + "DATABASE_URL_READ_REPLICA", +) +Surface = Literal["chat", "responses"] +TOOL_COUNT: Final = 300 + + +class EmbeddingDouble: + def __init__(self, mode: Literal["healthy", "held"], delay: float = 0.0) -> None: + self.release = threading.Event() + self.mode = mode + self.delay = delay + self.arrived: list[tuple[tuple[str, ...], float]] = [] + self.answered: list[tuple[tuple[str, ...], float]] = [] + + def document_inputs(self, descriptions: frozenset[str]) -> list[str]: + """Embedding inputs that are tool descriptions; query embeddings ('test', prompts) excluded.""" + return [text for inputs, _ in self.arrived for text in inputs if text in descriptions] + + def respond(self, request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=b'{"object": "list", "data": []}') + assert request.method == "POST" and request.target.endswith("/embeddings"), request.target + body: Final = json.loads(request.body) + raw: Final = body.get("input", []) + inputs: Final = tuple(str(text) for text in raw) if isinstance(raw, list) else (str(raw),) + self.arrived.append((inputs, time.monotonic())) + if self.mode == "held": + self.release.wait(timeout=120) + if self.delay: + time.sleep(self.delay) + data: Final = [ + { + "object": "embedding", + "index": index, + "embedding": [1.0, 0.0] if "invoice" in text.lower() else [0.0, 1.0], + } + for index, text in enumerate(inputs) + ] + self.answered.append((inputs, time.monotonic())) + return Reply( + body=json.dumps( + { + "object": "list", + "data": data, + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ).encode() + ) + + +def _chat_double() -> tuple[Callable[[Request], Reply], list[dict[str, object]]]: + bodies: list[dict[str, object]] = [] + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=b'{"object": "list", "data": []}') + body: Final = json.loads(request.body) + bodies.append(body) + identity: Final = uuid.uuid4().hex[:12] + usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + if request.target.endswith("/chat/completions"): + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{identity}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "done"}, + } + ], + "usage": usage, + } + ).encode() + ) + assert request.target.endswith("/responses"), request.target + return Reply( + body=json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + return respond, bodies + + +def _tool_names(body: dict[str, object]) -> tuple[str, ...]: + tools: Final = body.get("tools") + if not isinstance(tools, list): + return () + return tuple( + str(tool["name"] if "name" in tool else tool["function"]["name"]) for tool in tools if isinstance(tool, dict) + ) + + +def _peer_names(bodies: list[dict[str, object]], alias: str) -> frozenset[str]: + prefix: Final = f"{alias}-" + return frozenset(name for body in bodies for name in _tool_names(body) if name.startswith(prefix)) + + +def _tools(invoice_description: str = "Refund an invoice for a customer") -> tuple[ScriptedTool, ...]: + refund: Final = ScriptedTool( + "refund_invoice", + lambda params: text_result("refunded"), + description=invoice_description, + ) + weather: Final = tuple( + ScriptedTool( + f"weather_{n}", + lambda params: text_result("sunny"), + description=f"Look up the weather forecast for city {n}", + ) + for n in range(TOOL_COUNT - 1) + ) + return (refund, *weather) + + +def _all_names(alias: str) -> frozenset[str]: + return frozenset((f"{alias}-refund_invoice", *(f"{alias}-weather_{n}" for n in range(TOOL_COUNT - 1)))) + + +def _descriptions(invoice_description: str = "Refund an invoice for a customer") -> frozenset[str]: + return frozenset(tool.description for tool in _tools(invoice_description)) + + +def _config( + directory: Path, + *, + alias: str, + peer_url: str, + embeddings_url: str, + chat_url: str, + redis: OwnedRedis | None, +) -> Path: + document: Final = { + "model_list": [ + { + "model_name": "semantic-embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": f"{embeddings_url}/v1", + "api_key": "synthetic-embed-key", + }, + }, + { + "model_name": "chat-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{chat_url}/v1", + "api_key": "synthetic-chat-key", + }, + }, + { + "model_name": "responses-model", + "litellm_params": { + "model": "openai/responses/gpt-4o-mini", + "api_base": f"{chat_url}/v1", + "api_key": "synthetic-chat-key", + }, + }, + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + **({"coordination_redis": {"host": redis.host, "port": redis.port}} if redis is not None else {}), + }, + "mcp_servers": {alias: {"transport": "http", "url": peer_url}}, + "litellm_settings": {"mcp_semantic_tool_filter": SEMANTIC_SETTINGS}, + } + path: Final = directory / f"semantic-filter-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.dump(document)) + return path + + +def _send(proxy: Gateway, surface: Surface) -> httpx.Response: + if surface == "chat": + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": "chat-model", "messages": [{"role": "user", "content": PROMPT}], "tools": [MCP_REF]}, + ) + return proxy.request( + "POST", + "/v1/responses", + {"model": "responses-model", "input": PROMPT, "tools": [MCP_REF]}, + ) + + +def _request_is_filtered(proxy: Gateway, surface: Surface, bodies: list[dict[str, object]], alias: str) -> bool: + try: + response: Final = _send(proxy, surface) + except httpx.TransportError: + return False + return response.status_code == 200 and _peer_names(bodies[-1:], alias) == frozenset((f"{alias}-refund_invoice",)) + + +@contextmanager +def _proxy_without_readiness_gate( + gateway: Gateway, directory: Path, overrides: dict[str, str], *, config: Path +) -> Iterator[OwnedProxy]: + """owned_proxy_process without the readiness wait, for builds that must stay unready.""" + root: Final = _proxy_root() + environment: Final = _proxy_environment(gateway, overrides, NO_REDIS_ENV) + output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output.mkdir(parents=True, exist_ok=True) + command: Final = ( + sys.executable, + "-m", + "integration._support.proxy", + "--config", + str(config), + "--host", + "127.0.0.1", + "--num_workers", + "1", + "--timeout_worker_healthcheck", + "5", + *DB_PUSH, + ) + launch: Final = _launch(command, root, environment, output) + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS + ) as client: + yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), launch.process, launch.log) + finally: + _stop_launch(launch) + + +def test_redis_backed_proxy_is_live_while_300_tool_index_builds(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + embeddings: Final = EmbeddingDouble("held") + respond, bodies = _chat_double() + with ( + owned_redis(tmp_path) as redis, + scripted_peer(*_tools()) as peer, + wire_server(respond) as chat_wire, + wire_server(embeddings.respond) as embedding_wire, + ): + try: + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as proxy: + liveliness: Final = proxy.gateway.client.get("/health/liveliness") + assert liveliness.status_code == 200, liveliness.text + response: Final = _send(proxy.gateway, "chat") + assert response.status_code == 200, response.text + assert _peer_names(bodies[-1:], alias) == _all_names(alias), _peer_names(bodies[-1:], alias) + embeddings.release.set() + eventually( + lambda: _request_is_filtered(proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + finally: + embeddings.release.set() + + +def test_second_proxy_on_warm_redis_embeds_no_descriptions(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + descriptions: Final = _descriptions() + with owned_redis(tmp_path) as redis, scripted_peer(*_tools()) as peer: + embeddings: Final = EmbeddingDouble("healthy") + respond, bodies = _chat_double() + second_respond, second_bodies = _chat_double() + with ( + wire_server(respond) as chat_wire, + wire_server(second_respond) as second_chat_wire, + wire_server(embeddings.respond) as embedding_wire, + ): + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as first_proxy: + eventually( + lambda: _request_is_filtered(first_proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + embedded_so_far: Final = embeddings.document_inputs(descriptions) + config = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=second_chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: + assert _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias) + assert embeddings.document_inputs(descriptions) == embedded_so_far + + +def test_two_proxies_cold_starting_together_embed_each_description_once(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + descriptions: Final = _descriptions() + embeddings: Final = EmbeddingDouble("healthy", delay=10.0) + respond_a, bodies_a = _chat_double() + respond_b, bodies_b = _chat_double() + with ( + owned_redis(tmp_path) as redis, + scripted_peer(*_tools()) as peer, + wire_server(respond_a) as chat_wire_a, + wire_server(respond_b) as chat_wire_b, + wire_server(embeddings.respond) as embedding_wire, + ): + config_a: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire_a.url, + redis=redis, + ) + config_b: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire_b.url, + redis=redis, + ) + with ( + owned_proxy_process(gateway, tmp_path, {}, config=config_a) as proxy_a, + owned_proxy_process(gateway, tmp_path, {}, config=config_b) as proxy_b, + ): + eventually( + lambda: _request_is_filtered(proxy_a.gateway, "chat", bodies_a, alias), + lambda filtered: filtered, + seconds=120, + ) + eventually( + lambda: _request_is_filtered(proxy_b.gateway, "chat", bodies_b, alias), + lambda filtered: filtered, + seconds=60, + ) + document_inputs: Final = Counter(embeddings.document_inputs(descriptions)) + assert document_inputs == Counter({description: 1 for description in descriptions}) + + +def test_edited_description_reembeds_only_that_tool(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + descriptions_v1: Final = _descriptions() + edited: Final = "Cancel and refund an invoice instantly" + with owned_redis(tmp_path) as redis: + embeddings: Final = EmbeddingDouble("healthy") + respond, bodies = _chat_double() + second_respond, second_bodies = _chat_double() + with wire_server(embeddings.respond) as embedding_wire: + with ( + scripted_peer(*_tools()) as peer_v1, + wire_server(respond) as chat_wire, + ): + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer_v1.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as first_proxy: + eventually( + lambda: _request_is_filtered(first_proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + with ( + scripted_peer(*_tools(edited)) as peer_v2, + wire_server(second_respond) as second_chat_wire, + ): + config = _config( + tmp_path, + alias=alias, + peer_url=peer_v2.url, + embeddings_url=embedding_wire.url, + chat_url=second_chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: + eventually( + lambda: _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + document_inputs: Final = Counter( + text for inputs, _ in embeddings.arrived for text in inputs if text not in ("test", PROMPT) + ) + assert document_inputs == Counter({**{d: 1 for d in descriptions_v1}, edited: 1}) + + +def test_without_redis_startup_waits_for_the_index(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + embeddings: Final = EmbeddingDouble("held") + respond, bodies = _chat_double() + with ( + scripted_peer(*_tools()) as peer, + wire_server(respond) as chat_wire, + wire_server(embeddings.respond) as embedding_wire, + ): + try: + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=None, + ) + with _proxy_without_readiness_gate(gateway, tmp_path, {}, config=config) as proxy: + + def liveliness() -> int | None: + try: + return proxy.gateway.client.get("/health/liveliness", timeout=2).status_code + except httpx.TransportError: + return None + + deadline: Final = time.monotonic() + 5 + while time.monotonic() < deadline: + assert liveliness() != 200, "blocking build should keep liveliness from 200 while held" + time.sleep(0.5) + embeddings.release.set() + eventually(liveliness, lambda status: status == 200, seconds=60) + eventually( + lambda: _request_is_filtered(proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + finally: + embeddings.release.set() diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py new file mode 100644 index 00000000000..a5df27c2520 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py @@ -0,0 +1,208 @@ +""" +Unit tests for CachingToolEncoder, the LiteLLMRouterEncoder that serves +document embeddings from the shared ToolVectorStore. +""" + +import asyncio +import sys +from unittest.mock import Mock + +import pytest + +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 + + +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, + ) + + +def _make_encoder(recorded_inputs, store, embedding_identity="text-embedding-3-small"): + from litellm.proxy._experimental.mcp_server.semantic_tool_encoder import CachingToolEncoder + + return CachingToolEncoder( + litellm_router_instance=_fake_embedding_router(recorded_inputs), + model_name="text-embedding-3-small", + score_threshold=0.3, + store=store, + embedding_identity=embedding_identity, + ) + + +@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_wrong_length_cached_vector_is_reembedded_and_overwritten(): + """ + A cached vector whose length disagrees with a live dimension probe is + treated as a miss: it is re-embedded, served at the right length, and the + store entry is overwritten so the stale shape cannot leak into the index. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + key = tool_vector_key("text-embedding-3-small", "get a ticket") + store = _InMemoryToolVectorStore({key: (9.9, 9.9, 9.9, 9.9)}) + recorded = [] + encoder = _make_encoder(recorded, store) + + vectors = await encoder.aencode_documents(["get a ticket"]) + + assert vectors == [[1.0, 0.0]] + assert [batch for batch in recorded if batch != ["test"]] == [["get a ticket"]] + assert store.vectors[key] == (1.0, 0.0) 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..ef593481c27 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 @@ -129,25 +129,17 @@ async def test_semantic_filter_basic_filtering(): ) # Assertions - validate filtering mechanics work - assert ( - len(filtered) <= 3 - ), f"Should return at most 3 tools (top_k), got {len(filtered)}" + assert len(filtered) <= 3, f"Should return at most 3 tools (top_k), got {len(filtered)}" assert len(filtered) > 0, "Should return at least some tools" - assert len(filtered) < len( - tools - ), f"Should filter down from {len(tools)} tools, got {len(filtered)}" + assert len(filtered) < len(tools), f"Should filter down from {len(tools)} tools, got {len(filtered)}" # Validate tools are actual MCPTool objects for tool in filtered: assert hasattr(tool, "name"), "Filtered result should be MCPTool with name" - assert hasattr( - tool, "description" - ), "Filtered result should be MCPTool with description" + assert hasattr(tool, "description"), "Filtered result should be MCPTool with description" filtered_names = [t.name for t in filtered] - print( - f"✅ Successfully filtered {len(tools)} tools down to top {len(filtered)}: {filtered_names}" - ) + print(f"✅ Successfully filtered {len(tools)} tools down to top {len(filtered)}: {filtered_names}") print(f" Filter respects top_k parameter correctly") @@ -226,12 +218,7 @@ async def test_semantic_filter_disabled(): SemanticMCPToolFilter, ) - tools = [ - MCPTool( - name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"} - ) - for i in range(10) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(10)] mock_router = Mock() @@ -251,9 +238,7 @@ async def test_semantic_filter_disabled(): ) # Should return all tools when disabled - assert len(filtered) == len( - tools - ), f"Expected all {len(tools)} tools, got {len(filtered)}" + assert len(filtered) == len(tools), f"Expected all {len(tools)} tools, got {len(filtered)}" @pytest.mark.asyncio @@ -373,12 +358,7 @@ async def test_semantic_filter_hook_triggers_on_completion(): ) # Prepare data - completion request with tools - tools = [ - MCPTool( - name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"} - ) - for i in range(10) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(10)] # Build router with the tools before filtering filter_instance._build_router(tools) @@ -408,9 +388,7 @@ async def test_semantic_filter_hook_triggers_on_completion(): # Assertions assert result is not None, "Hook should return modified data" assert "tools" in result, "Result should contain tools" - assert len(result["tools"]) < len( - tools - ), f"Hook should filter tools, got {len(result['tools'])}/{len(tools)}" + assert len(result["tools"]) < len(tools), f"Hook should filter tools, got {len(result['tools'])}/{len(tools)}" print(f"✅ Hook triggered correctly: {len(tools)} -> {len(result['tools'])} tools") @@ -557,18 +535,12 @@ async def test_semantic_filter_hook_preserves_native_tools(): filtered = result["tools"] # Native tools must survive - native_in_result = [ - t for t in filtered if isinstance(t, dict) and t.get("type") == "function" - ] - assert ( - len(native_in_result) == 2 - ), f"Both native tools must survive, got {len(native_in_result)}" + native_in_result = [t for t in filtered if isinstance(t, dict) and t.get("type") == "function"] + assert len(native_in_result) == 2, f"Both native tools must survive, got {len(native_in_result)}" # MCP tools should be filtered (top_k=2) mcp_in_result = [t for t in filtered if not isinstance(t, dict)] - assert ( - len(mcp_in_result) <= 2 - ), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" + assert len(mcp_in_result) <= 2, f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" # Total should be native + filtered MCP assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}" @@ -579,12 +551,8 @@ async def test_semantic_filter_hook_preserves_native_tools(): # Stats should report MCP-only counts, not inflated with native tools stats = result["metadata"]["litellm_semantic_filter_stats"] mcp_before, mcp_after = stats.split("->") - assert ( - int(mcp_before) == 5 - ), f"Stats 'from' should be MCP count (5), got {mcp_before}" - assert int(mcp_after) == len( - mcp_in_result - ), f"Stats 'to' should match filtered MCP count, got {mcp_after}" + assert int(mcp_before) == 5, f"Stats 'from' should be MCP count (5), got {mcp_before}" + assert int(mcp_after) == len(mcp_in_result), f"Stats 'to' should match filtered MCP count, got {mcp_after}" print( f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} " @@ -679,20 +647,15 @@ async def test_semantic_filter_hook_all_native_tools(): filtered = result["tools"] # All native tools must pass through - assert ( - len(filtered) == 3 - ), f"All 3 native tools must pass through, got {len(filtered)}" + assert len(filtered) == 3, f"All 3 native tools must pass through, got {len(filtered)}" # No spurious semantic filter stats (P2 fix) - assert ( - "litellm_semantic_filter_stats" not in result["metadata"] - ), "Should NOT emit semantic filter stats for all-native-tool requests" - - print( - f"✅ Hook passes through all {len(filtered)} native tools, " - f"no spurious filter headers emitted" + assert "litellm_semantic_filter_stats" not in result["metadata"], ( + "Should NOT emit semantic filter stats for all-native-tool requests" ) + print(f"✅ Hook passes through all {len(filtered)} native tools, no spurious filter headers emitted") + @requires_semantic_router @pytest.mark.asyncio @@ -758,8 +721,7 @@ async def test_semantic_filter_hook_responses_api_name_collision(): # Verify classification: should be native, not MCP assert not hook._is_mcp_tool(responses_api_tool), ( - "Responses API tool with type=function + top-level name " - "should be classified as native, not MCP" + "Responses API tool with type=function + top-level name should be classified as native, not MCP" ) # Full hook test: all-native request should preserve tools @@ -890,9 +852,9 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools(): for name in allowed_tools: assert name in expanded_names, "Selected tool names must come from the expanded tools" - assert ( - "litellm_semantic_filter_stats" in result["metadata"] - ), "Filter stats must be emitted for the litellm_proxy expansion path" + assert "litellm_semantic_filter_stats" in result["metadata"], ( + "Filter stats must be emitted for the litellm_proxy expansion path" + ) stats = result["metadata"]["litellm_semantic_filter_stats"] total, selected = stats.split("->") assert int(total) == 5, f"Stats 'from' should be pre-filter expanded count (5), got {total}" @@ -1274,9 +1236,9 @@ async def test_semantic_filter_hook_expansion_skips_filter_when_disabled(): "require_approval": "never", } ], "The MCP reference must be left intact for the MCP gateway to expand" - assert ( - "litellm_semantic_filter_stats" not in data["metadata"] - ), "No filter stats may be emitted when the filter is disabled" + assert "litellm_semantic_filter_stats" not in data["metadata"], ( + "No filter stats may be emitted when the filter is disabled" + ) print("✅ Disabled filter: MCP reference untouched, no spurious stats") @@ -1380,9 +1342,7 @@ async def test_semantic_filter_hook_preserves_tool_order(): assert filtered[1] is native_tool, "Second tool should be native_tool" assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b" - print( - "✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering" - ) + print("✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering") class TestGetToolsByNames: @@ -1492,9 +1452,7 @@ class TestGetToolsByNames: filter_instance = self._make_filter() available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}] - matched = filter_instance._get_tools_by_names( - ["fs-read_file", "api-fs-read_file"], available_tools - ) + matched = filter_instance._get_tools_by_names(["fs-read_file", "api-fs-read_file"], available_tools) assert len(matched) == 1 @@ -1525,9 +1483,7 @@ class TestGetToolsByNames: {"name": "litellm_fs-delete", "description": "delete"}, ] - matched = filter_instance._get_tools_by_names( - ["fs-write", "fs-delete", "fs-read"], available_tools - ) + matched = filter_instance._get_tools_by_names(["fs-write", "fs-delete", "fs-read"], available_tools) names = [t["name"] for t in matched] assert names == [ @@ -1682,10 +1638,7 @@ async def test_semantic_filter_fails_closed_on_query_time_context_window_error() state = {"raise_context_error": False} filter_instance = _make_context_window_filter(state) - tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(5) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(5)] filter_instance._build_router(tools) assert filter_instance.tool_router is not None @@ -1715,10 +1668,7 @@ async def test_semantic_filter_records_build_time_context_window_error(): state = {"raise_context_error": True} filter_instance = _make_context_window_filter(state) - tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(5) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(5)] filter_instance._build_router(tools) assert filter_instance.tool_router is None @@ -1749,10 +1699,7 @@ async def test_semantic_filter_hook_fails_closed_on_context_window_error(): state = {"raise_context_error": False} filter_instance = _make_context_window_filter(state) - tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(5) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(5)] filter_instance._build_router(tools) hook = SemanticToolFilterHook(filter_instance) @@ -1861,10 +1808,7 @@ async def test_semantic_filter_hook_ignores_build_error_for_native_only_tools(): state = {"raise_context_error": True} filter_instance = _make_context_window_filter(state) - mcp_tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(3) - ] + mcp_tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(3)] filter_instance._build_router(mcp_tools) assert filter_instance.context_window_error is not None @@ -1921,9 +1865,7 @@ def test_is_context_window_error_detection_variants(): def wrap_without_explicit_chaining(): try: - raise litellm.ContextWindowExceededError( - message="overflow", model="m", llm_provider="openai" - ) + raise litellm.ContextWindowExceededError(message="overflow", model="m", llm_provider="openai") except litellm.ContextWindowExceededError: raise ValueError("wrapper without explicit chaining") @@ -2239,3 +2181,288 @@ 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], + reindex_changed_descriptions=True, + ) + + 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_request_body_tool_cannot_trigger_description_reindex(): + """ + A caller-supplied tool with a registered name but a made-up description + must not re-embed or replace the worker-wide index entry. Only the + registry-expanded path is trusted to reindex changed descriptions. + """ + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + 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"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + spoofed = { + "name": "srv-a-get_thing", + "description": "totally made up description", + "inputSchema": {"type": "object"}, + } + hook = SemanticToolFilterHook(filter_instance) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "weather"}], + "tools": [spoofed], + "metadata": {}, + } + + await hook.async_pre_call_hook(user_api_key_dict=Mock(), cache=Mock(), data=data, call_type="completion") + + assert recorded_inputs == [["weather"]], "request-body tools must not trigger document embeddings" + indexed_tool = filter_instance._tool_map["srv-a-get_thing"] + assert filter_instance.extract_tool_info(indexed_tool)[1] == "get the weather for a city" + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_hook_expanded_path_reindexes_changed_description(): + """ + Tools the proxy expanded from the MCP registry are trusted: a changed + description re-embeds exactly that one tool through _filter_expanded_tools. + """ + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + 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"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + edited = { + "name": "srv-a-get_thing", + "description": "get a Linear ticket by id", + "inputSchema": {"type": "object"}, + } + hook = SemanticToolFilterHook(filter_instance) + filtered = await hook._filter_expanded_tools( + data={"messages": [{"role": "user", "content": "linear ticket"}]}, + expanded_tools=[edited], + ) + + assert [edited] == filtered + assert recorded_inputs == [["get a Linear ticket by id"], ["linear ticket"]] + 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_failed_reindex_retries_on_next_request(): + """ + When reindexing a changed description raises mid-embed, the tool must stay + 'missing' so the next request retries it instead of silently dropping it + from matching (its route was already deleted from the router). + """ + 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"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + original_aembedding = filter_instance.router_instance.aembedding + calls = {"n": 0} + + async def fail_once(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("embedding backend down") + return await original_aembedding(*args, **kwargs) + + filter_instance.router_instance.aembedding = fail_once + + edited = MCPTool( + name="srv-a-get_thing", + description="get a Linear ticket by id", + inputSchema={"type": "object"}, + ) + await filter_instance.filter_tools( + query="linear ticket", + available_tools=[edited], + reindex_changed_descriptions=True, + ) + assert "srv-a-get_thing" not in filter_instance._tool_map, ( + "failed re-embed must leave the tool missing, not half-indexed" + ) + + filtered = await filter_instance.filter_tools( + query="linear ticket", + available_tools=[edited], + reindex_changed_descriptions=True, + ) + assert [edited] == filtered + assert any("get a Linear ticket by id" in text for batch in recorded_inputs for text in batch) + assert "srv-a-get_thing" in filter_instance._tool_map + + +@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 + + +def test_embedding_identity_marks_deployment_dimensions(): + """ + The shared-vector key identity must change when a deployment's embedding + dimensions change, and stay byte-identical to the pre-dimensions format + when no deployment sets them (warm caches must keep hitting). + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + + router_without = Mock() + router_without.get_model_list = lambda model_name: [{"litellm_params": {"model": "amazon.titan-embed-text-v2:0"}}] + router_with = Mock() + router_with.get_model_list = lambda model_name: [ + {"litellm_params": {"model": "amazon.titan-embed-text-v2:0", "dimensions": 512}} + ] + router_with_api_base = Mock() + router_with_api_base.get_model_list = lambda model_name: [ + { + "litellm_params": { + "model": "amazon.titan-embed-text-v2:0", + "dimensions": 512, + "api_base": "https://bedrock-endpoint.example.com", + } + } + ] + + base = SemanticMCPToolFilter( + embedding_model="titan-alias", + litellm_router_instance=router_without, + enabled=True, + ) + sized = SemanticMCPToolFilter( + embedding_model="titan-alias", + litellm_router_instance=router_with, + enabled=True, + ) + based = SemanticMCPToolFilter( + embedding_model="titan-alias", + litellm_router_instance=router_with_api_base, + enabled=True, + ) + + assert base.embedding_identity == "titan-alias|amazon.titan-embed-text-v2:0" + assert sized.embedding_identity == "titan-alias|amazon.titan-embed-text-v2:0@dimensions=512" + assert ( + based.embedding_identity + == "titan-alias|amazon.titan-embed-text-v2:0@dimensions=512@api_base=https://bedrock-endpoint.example.com" + ) 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..b53bab95e9d --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py @@ -0,0 +1,366 @@ +""" +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 asyncio +import os +import subprocess +import sys +from unittest.mock import AsyncMock, Mock, patch + +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 + self.block_put_many = False + self.put_started = asyncio.Event() + self.put_block = asyncio.Event() + + 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): + if self.block_put_many: + self.put_started.set() + await self.put_block.wait() + 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 _ReadOnlyRedisCache: + """Reads work; every write reports failure like a read-only replica.""" + + def __init__(self): + self.backing = {} + + 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): + raise ConnectionError("NOPERM read only replica") + + async def async_get_cache(self, key, parent_otel_span=None, **kwargs): + return self.backing.get(key) + + async def async_set_cache(self, key, value, **kwargs): + return False + + +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", lock_wait_timeout_s=300.0): + 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, + lock_wait_timeout_s=lock_wait_timeout_s, + ) + + +@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_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"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_full_build_reads_the_store_a_constant_number_of_times(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tools = [_tool(f"srv-a_get_tool_{i}", f"do thing {i}") for i in range(20)] + store = _InMemoryToolVectorStore() + filter_instance = _make_filter([], store=store) + + 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=tools)), + ): + await filter_instance.build_router_from_mcp_registry() + + assert filter_instance.tool_router is not None + assert store.get_many_calls <= 3, ( + f"{store.get_many_calls} get_many calls for {len(tools)} tools; reads must not scale with the tool count" + ) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_cancelling_a_build_while_holding_the_lock_releases_it(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + store.block_put_many = True + filter_instance = _make_filter([], store=store) + + build_task = asyncio.create_task(filter_instance._abuild_router(tools)) + await asyncio.wait_for(store.put_started.wait(), timeout=5) + assert store.acquire_calls == 1 + + build_task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(build_task, timeout=5) + assert store.release_calls == 1 + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_wait_times_out_and_the_build_falls_back_to_local_embed(): + """ + A store that never fills and never grants the lock must not wait forever: + after lock_wait_timeout_s the build embeds locally, the startup index + flips ready, and the filter still selects tools. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + store.lock_result = False + recorded = [] + filter_instance = _make_filter(recorded, store=store, lock_wait_timeout_s=0.05) + filter_instance._lock_poll_interval_s = 0.01 + + 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=tools)), + ): + await filter_instance.build_router_from_mcp_registry() + + assert filter_instance.startup_index_ready is True + assert filter_instance.tool_router is not None + assert store.release_calls == 0 + assert [batch for batch in recorded if batch != ["test"]] == [["get a ticket", "weather forecast"]] + filtered = await filter_instance.filter_tools( + query="refund a ticket", available_tools=tools, reindex_changed_descriptions=True + ) + assert filtered == [tools[0]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_acquire_embeds_locally_when_writes_fail_and_no_pod_holds_the_lock(): + 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 = _ReadOnlyRedisCache() + store = RedisToolVectorStore(redis, PodLockManager(redis)) + + assert await store.try_acquire_build_lock() is True + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_acquire_returns_false_when_another_pod_holds_the_lock(): + 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 = _ReadOnlyRedisCache() + lock_manager = PodLockManager(redis) + redis.backing[PodLockManager.get_redis_lock_key("mcp_semantic_tool_index_build")] = "other-pod" + store = RedisToolVectorStore(redis, lock_manager) + + assert await store.try_acquire_build_lock() is False + assert store._degraded is False + + +def test_store_and_hook_import_without_semantic_router(): + """semantic_router is an optional extra; the store and the hook must import when it is absent.""" + root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) + ) + result = subprocess.run( + [ + sys.executable, + "-I", + "-c", + "import sys; sys.path.insert(0, " + repr(root) + "); " + "sys.modules['semantic_router'] = None; " + "import litellm.proxy._experimental.mcp_server.semantic_tool_index_store; " + "import litellm.proxy.hooks.mcp_semantic_filter", + ], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index a1e22780d57..5188ead54e7 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -96,7 +96,9 @@ def test_proxy_config_upserts_existing_deployment_without_duplicating_it(monkeyp @pytest.mark.asyncio -async def test_proxy_config_preserves_models_when_config_read_fails(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +async def test_proxy_config_preserves_models_when_config_read_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: from litellm.proxy import proxy_server router = litellm.Router(model_list=[_deployment("config-deployment").to_json(exclude_none=True)]) @@ -129,6 +131,8 @@ async def test_proxy_config_keeps_database_models_and_deletes_stale_models( assert desired == frozenset({"kept-deployment"}) assert router.get_model_ids() == ["kept-deployment"] + + from tests._master_key import MASTER_KEY @@ -144,7 +148,6 @@ async def test_proxy_config_loads_lens_store_from_yaml(tmp_path, monkeypatch) -> assert is_lens_tracing_enabled(settings["tracing"], {}) is True - @pytest.mark.asyncio @pytest.mark.parametrize("shutdown_error", [False, True]) async def test_tracing_config_automatically_exports_spend_without_a_storage_dependency( @@ -170,7 +173,9 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe outcome: Final = pytest.raises(RuntimeError, match="shutdown failure") if shutdown_error else nullcontext() with outcome: async with manage_tracing(enabled=True, client_factory=client): - logger: Final = next(callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter)) + logger: Final = next( + callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter) + ) await logger.async_log_success_event( {"standard_logging_object": {"id": "response-1", "response_cost": 0.25}}, None, None, None ) @@ -186,7 +191,6 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe assert logger.task is not None and logger.task.done() and not logger.task.cancelled() - # --------------------------------------------------------------------------- # _is_remote_module_url # --------------------------------------------------------------------------- @@ -4835,7 +4839,9 @@ async def test_ProxyConfig__update_config_from_db_keeps_keys_the_config_file_omi @pytest.mark.asyncio @pytest.mark.parametrize("ui_settings_already_synced", [False, True]) -async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints(monkeypatch, ui_settings_already_synced): +async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints( + monkeypatch, ui_settings_already_synced +): from litellm.proxy import proxy_server pc = ProxyConfig() @@ -5607,3 +5613,103 @@ 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 + + +@pytest.mark.asyncio +async def test_proxy_config_semantic_filter_db_sync_reinits_after_finished_failed_build(monkeypatch): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + settings: Final = { + "enabled": True, + "embedding_model": "embed", + "top_k": 3, + "similarity_threshold": 0.5, + } + build_task: Final = asyncio.create_task(asyncio.sleep(0)) + await build_task + 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 + + await pc._init_semantic_filter_settings_in_db(prisma_client=MagicMock()) + + removed.assert_called_once() + + +@pytest.mark.asyncio +async def test_proxy_config_semantic_filter_db_sync_cancels_orphaned_index_build(monkeypatch): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + settings: Final = { + "enabled": True, + "embedding_model": "embed", + "top_k": 3, + "similarity_threshold": 0.5, + } + changed_settings: Final = {**settings, "top_k": 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": changed_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()) + + await asyncio.sleep(0) + assert build_task.cancelled() + removed.assert_called_once() + finally: + gate.set() + if not build_task.done(): + build_task.cancel()