From 74b4eab364348aa8cb22e4e15060ba78e083bd42 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sun, 3 May 2026 10:27:53 +0000 Subject: [PATCH] fix(vector_store): cache use-time embedding-config resolution Hold the resolved config in a process-memory TTL cache so the request-handling path doesn't run litellm_proxymodeltable.find_first on every vector-store call. --- .../management_endpoints.py | 35 +++++++++++++ .../test_vector_store_endpoints.py | 50 +++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index bff43cacaa..cbb3d92718 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -16,6 +16,7 @@ from fastapi import APIRouter, Depends, HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker @@ -45,6 +46,28 @@ _LITELLM_PARAMS_MASKER = SensitiveDataMasker() _REDACT_LITELLM_PARAMS_MAX_DEPTH = 10 +# Use-time embedding-config resolution runs on every vector-store request +# whose persisted row carries only a model reference (the post-fix shape). +# Without a cache, that's one ``litellm_proxymodeltable.find_first`` per +# request — the no-DB-in-critical-path rule. Hold the resolved config in +# memory for a short TTL so a hot model name pays the DB lookup at most +# once per ``_EMBEDDING_CONFIG_CACHE_TTL`` seconds. Cleartext credentials +# only ever live in process memory (never persisted, never echoed in +# management responses), so the cache doesn't widen the disclosure surface. +_EMBEDDING_CONFIG_CACHE_TTL = 60 +_EMBEDDING_CONFIG_CACHE_MAX_SIZE = 256 +_embedding_config_cache: Optional[InMemoryCache] = None + + +def _get_embedding_config_cache() -> InMemoryCache: + global _embedding_config_cache + if _embedding_config_cache is None: + _embedding_config_cache = InMemoryCache( + max_size_in_memory=_EMBEDDING_CONFIG_CACHE_MAX_SIZE, + default_ttl=_EMBEDDING_CONFIG_CACHE_TTL, + ) + return _embedding_config_cache + def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> Any: """ @@ -303,6 +326,11 @@ async def _resolve_embedding_config( This function first checks the router for config-defined models, then falls back to the database. This allows users to use models defined in either location. + Results are cached in process memory for ``_EMBEDDING_CONFIG_CACHE_TTL`` + seconds so the request-handling path doesn't hit the database on every + vector-store call. Negative results (model not found) are intentionally + not cached to avoid blocking a freshly-added model behind the TTL. + Args: embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") prisma_client: The Prisma client instance @@ -314,6 +342,11 @@ async def _resolve_embedding_config( if not embedding_model: return None + cache = _get_embedding_config_cache() + cached = cache.get_cache(embedding_model) + if cached is not None: + return cached + # Import llm_router if not provided if llm_router is None: try: @@ -330,6 +363,7 @@ async def _resolve_embedding_config( verbose_proxy_logger.debug( f"Resolved embedding config from router for model {embedding_model}" ) + cache.set_cache(embedding_model, router_config) return router_config # Fall back to database @@ -341,6 +375,7 @@ async def _resolve_embedding_config( verbose_proxy_logger.debug( f"Resolved embedding config from database for model {embedding_model}" ) + cache.set_cache(embedding_model, db_config) return db_config verbose_proxy_logger.debug( diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 8fc7bc4e26..81b67e8bc5 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -50,6 +50,19 @@ def _serialize_litellm_params(litellm_params): return json.dumps(litellm_params or {}) +@pytest.fixture(autouse=True) +def _reset_embedding_config_cache(): + """The use-time embedding-config resolver caches results in process + memory across calls. Reset it before every test so the resolver + actually exercises the router/DB path under test instead of returning + a value cached by an earlier test.""" + from litellm.proxy.vector_store_endpoints import management_endpoints + + management_endpoints._embedding_config_cache = None + yield + management_endpoints._embedding_config_cache = None + + @pytest.mark.asyncio async def test_router_avector_store_search_passes_correct_args(): """ @@ -1684,6 +1697,43 @@ async def test_resolve_embedding_config_tries_router_then_db(): mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_not_called() +@pytest.mark.asyncio +async def test_resolve_embedding_config_caches_result(): + """The first lookup should hit the router/DB; subsequent lookups for + the same model name should return the cached value without touching + the router or the database.""" + from litellm.types.router import Deployment, LiteLLM_Params + + mock_prisma_client = MagicMock() + mock_router = MagicMock() + + mock_litellm_params = MagicMock(spec=LiteLLM_Params) + mock_litellm_params.api_key = "router-api-key" + mock_litellm_params.api_base = "https://router-api-base.com" + mock_litellm_params.api_version = None + + mock_deployment = MagicMock(spec=Deployment) + mock_deployment.litellm_params = mock_litellm_params + mock_router.get_deployment_by_model_group_name.return_value = mock_deployment + + first = await _resolve_embedding_config( + embedding_model="cached-model", + prisma_client=mock_prisma_client, + llm_router=mock_router, + ) + assert first is not None + assert mock_router.get_deployment_by_model_group_name.call_count == 1 + + second = await _resolve_embedding_config( + embedding_model="cached-model", + prisma_client=mock_prisma_client, + llm_router=mock_router, + ) + assert second == first + # Router (and by extension the DB) was not consulted again. + assert mock_router.get_deployment_by_model_group_name.call_count == 1 + + @pytest.mark.asyncio async def test_resolve_embedding_config_falls_back_to_db(): """Test that _resolve_embedding_config falls back to DB when router doesn't have the model."""