diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 86e316e7f40..ccf15c206b0 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -8,6 +8,9 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import jsonify_object +from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _resolve_embedding_config, +) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store, get_litellm_managed_vector_store, @@ -56,6 +59,30 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( if "litellm_params" in vector_store_to_run: litellm_params = vector_store_to_run.get("litellm_params", {}) or {} + # Resolve ``litellm_embedding_config`` here, at request-handling + # time, instead of at row-creation time. The resolved + # ``api_key`` / ``api_base`` / ``api_version`` lives only in + # this per-request ``data`` dict and is never persisted. + # Legacy rows that already carry a resolved (cleartext) + # ``litellm_embedding_config`` skip the lookup and pass through + # unchanged so the embed call keeps working. + embedding_model = litellm_params.get("litellm_embedding_model") + if embedding_model and not litellm_params.get("litellm_embedding_config"): + from litellm.proxy.proxy_server import prisma_client + + resolved_config = await _resolve_embedding_config( + embedding_model=embedding_model, prisma_client=prisma_client + ) + if resolved_config: + # Build a fresh dict via spread instead of mutating + # ``litellm_params`` in place — the registry hands back + # a reference to its cached object, so an in-place + # update would persist the resolved cleartext into the + # in-memory cache for the lifetime of the process. + litellm_params = { + **litellm_params, + "litellm_embedding_config": resolved_config, + } data.update(litellm_params) return data diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 99a2085bfcd..cbb3d927184 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( @@ -432,20 +467,17 @@ async def create_vector_store_in_db( if user_id is not None: data_to_create["user_id"] = user_id - # Handle litellm_params - always provide at least an empty dict + # Handle litellm_params - always provide at least an empty dict. + # The earlier behaviour resolved ``litellm_embedding_config`` from the + # admin-configured router/DB model and persisted the cleartext result + # (``api_key``, ``api_base``, ``api_version``) into this row. That + # exposed every env-stored embedding-model credential on the + # ``/vector_store/{new,info,update,list}`` responses. Keep the user's + # raw ``litellm_embedding_model`` reference; resolution now happens in + # ``_update_request_data_with_litellm_managed_vector_store_registry`` + # at request-handling time so the cleartext config exists only in + # per-request memory and never reaches the database. if litellm_params: - # Auto-resolve embedding config if embedding model is provided but config is not - embedding_model = litellm_params.get("litellm_embedding_model") - if embedding_model and not litellm_params.get("litellm_embedding_config"): - resolved_config = await _resolve_embedding_config( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if resolved_config: - litellm_params["litellm_embedding_config"] = resolved_config - verbose_proxy_logger.info( - f"Auto-resolved embedding config for model {embedding_model}" - ) - litellm_params_dict = GenericLiteLLMParams(**litellm_params).model_dump( exclude_none=True ) @@ -531,10 +563,19 @@ async def new_vector_store( user_id=user_api_key_dict.user_id, ) + # Apply the same litellm_params redaction the list / info / update + # endpoints already use, so a caller-supplied credential or a + # cleartext value persisted by an earlier proxy version doesn't + # come back in the response. + response_vs = LiteLLM_ManagedVectorStore(**new_vector_store) + response_vs["litellm_params"] = _redact_sensitive_litellm_params( + new_vector_store.get("litellm_params") + ) + return { "status": "success", "message": f"Vector store {vector_store.get('vector_store_id')} created successfully", - "vector_store": new_vector_store, + "vector_store": response_vs, } except Exception as e: verbose_proxy_logger.exception(f"Error creating vector store: {str(e)}") @@ -865,24 +906,15 @@ async def update_vector_store( update_data["vector_store_metadata"] ) - # Handle litellm_params if provided + # Handle litellm_params if provided. As with the create path, the + # embedding-config auto-resolve previously persisted cleartext + # credentials into the row; resolution now happens at request- + # handling time in + # ``_update_request_data_with_litellm_managed_vector_store_registry`` + # so this row only ever stores the user-supplied + # ``litellm_embedding_model`` reference. if "litellm_params" in update_data: _input_litellm_params: dict = update_data.get("litellm_params", {}) or {} - - # Auto-resolve embedding config if embedding model is provided but config is not - embedding_model = _input_litellm_params.get("litellm_embedding_model") - if embedding_model and not _input_litellm_params.get( - "litellm_embedding_config" - ): - resolved_config = await _resolve_embedding_config( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if resolved_config: - _input_litellm_params["litellm_embedding_config"] = resolved_config - verbose_proxy_logger.info( - f"Auto-resolved embedding config for model {embedding_model}" - ) - litellm_params_dict = GenericLiteLLMParams( **_input_litellm_params ).model_dump(exclude_none=True) 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 e67a04c749a..81b67e8bc50 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 @@ -36,6 +36,33 @@ from litellm.proxy.vector_store_endpoints.utils import ( from litellm.types.utils import LlmProviders +def _serialize_litellm_params(litellm_params): + """Serialize ``litellm_params`` to a string for substring assertions. + + The redact helper preserves the persisted shape — string in, string + out; dict in, dict out — so callers that just want to assert "this + secret never appears" need a single text representation either way. + """ + import json + + if isinstance(litellm_params, str): + return 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(): """ @@ -170,6 +197,93 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry(): assert result == original_data +@pytest.mark.asyncio +async def test_update_request_data_resolves_embedding_config_at_use_time(): + """When the persisted vector store row carries only a + ``litellm_embedding_model`` reference (the new behaviour after + moving the auto-resolve out of write time), the request-handling + layer must resolve the embedding config so the downstream embed + call still has ``api_key`` / ``api_base`` / ``api_version``. The + resolved config lives in this per-request data dict only — never + persisted.""" + mock_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "test_store", + "custom_llm_provider": "azure_ai", + "litellm_params": { + "litellm_embedding_model": "azure/text-embedding-3-large", + # Note: no litellm_embedding_config persisted + }, + } + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( + mock_vector_store + ) + + resolved = { + "api_key": "use-time-resolved-key", + "api_base": "https://my-azure.example", + "api_version": "2024-09-01", + } + + with ( + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", + new=AsyncMock(return_value=resolved), + ), + ): + result = await _update_request_data_with_litellm_managed_vector_store_registry( + data={}, vector_store_id="test_store" + ) + + assert result["litellm_embedding_model"] == "azure/text-embedding-3-large" + assert result["litellm_embedding_config"] == resolved + + +@pytest.mark.asyncio +async def test_update_request_data_passes_through_legacy_embedding_config(): + """A vector store row created by an older proxy version may already + carry a fully-resolved ``litellm_embedding_config`` in its persisted + ``litellm_params`` (the very leak this PR closes). Those legacy rows + must still work — the use-time resolver skips re-resolution when + the config is already present so the embed call keeps succeeding.""" + legacy_config = { + "api_key": "legacy-cleartext-key", + "api_base": "https://legacy-azure.example", + "api_version": "2024-01-01", + } + mock_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "legacy_store", + "custom_llm_provider": "azure_ai", + "litellm_params": { + "litellm_embedding_model": "azure/text-embedding-3-large", + "litellm_embedding_config": legacy_config, + }, + } + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( + mock_vector_store + ) + + resolve_mock = AsyncMock() + + with ( + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", + new=resolve_mock, + ), + ): + result = await _update_request_data_with_litellm_managed_vector_store_registry( + data={}, vector_store_id="legacy_store" + ) + + assert result["litellm_embedding_config"] == legacy_config + resolve_mock.assert_not_awaited() + + class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function.""" @@ -1417,20 +1531,25 @@ async def test_new_vector_store_auto_resolves_embedding_config(): ) assert result["status"] == "success" - # Verify that embedding config was resolved and included in the create call + # Auto-resolve no longer happens at create time — the persisted row + # carries only the model reference, never the resolved cleartext + # credential. Resolution now happens at request-handling time inside + # ``_update_request_data_with_litellm_managed_vector_store_registry``, + # where the resolved config lives in per-request memory and is never + # written to the database. litellm_params_json = captured_create_data.get("litellm_params") assert litellm_params_json is not None litellm_params_dict = json.loads(litellm_params_json) - assert "litellm_embedding_config" in litellm_params_dict - assert ( - litellm_params_dict["litellm_embedding_config"]["api_key"] == "resolved-api-key" - ) - assert ( - litellm_params_dict["litellm_embedding_config"]["api_base"] - == "https://api.openai.com" - ) - assert ( - litellm_params_dict["litellm_embedding_config"]["api_version"] == "2024-01-01" + assert "litellm_embedding_config" not in litellm_params_dict + assert litellm_params_dict["litellm_embedding_model"] == "text-embedding-ada-002" + + # The response must also not echo a cleartext credential — even on + # the create response, where redaction guards against caller-supplied + # cleartext or pre-existing rows that were created by an earlier + # proxy version. + response_vs = result["vector_store"] + assert "resolved-api-key" not in _serialize_litellm_params( + response_vs.get("litellm_params") ) @@ -1578,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.""" @@ -1687,21 +1843,18 @@ async def test_new_vector_store_auto_resolves_from_router(): ) assert result["status"] == "success" - # Verify that embedding config was resolved from router and included in the create call + # Resolution against the router happens at request-handling time now, + # not at row creation. The persisted ``litellm_params`` carries only + # the model reference, never the cleartext credential. litellm_params_json = captured_create_data.get("litellm_params") assert litellm_params_json is not None litellm_params_dict = json.loads(litellm_params_json) - assert "litellm_embedding_config" in litellm_params_dict - assert ( - litellm_params_dict["litellm_embedding_config"]["api_key"] - == "router-resolved-api-key" - ) - assert ( - litellm_params_dict["litellm_embedding_config"]["api_base"] - == "https://router-resolved-base.com" - ) - assert ( - litellm_params_dict["litellm_embedding_config"]["api_version"] == "2024-03-01" + assert "litellm_embedding_config" not in litellm_params_dict + assert litellm_params_dict["litellm_embedding_model"] == "config-embedding-model" + + response_vs = result["vector_store"] + assert "router-resolved-api-key" not in _serialize_litellm_params( + response_vs.get("litellm_params") )