scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

Merge pull request #27082 from stuxf/fix/vector-store-cred-leak

Browse source 
fix(vector_store): resolve embedding config at request time, never persist creds
This commit is contained in:
yuneng-jiang 2026-05-04 15:55:40 -07:00 • committed by GitHub
commit 0c0b5e005f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 265 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/proxy/vector_store_endpoints/endpoints.py
View file

@ -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

92
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -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)

199
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -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")
)