scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(rag): store-wins merge, single lookup, allowlisted search params

Browse source 
rag_query reuses the store resolved during authorization instead of a
second registry lookup, merges registry data store-wins so callers
cannot override a managed store's provider or credentials, and logs ids
instead of the merged config, which can carry resolved credentials.
aquery forwards only allowlisted retrieval_config keys to vector store
search, keeping caller-supplied connection overrides like api_base and
api_key away from the search call
This commit is contained in:
mateo-berri 2026-09-01 12:52:45 -07:00
parent 3914de24ef
commit babe7816ad
6 changed files with 149 additions and 84 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
litellm/proxy/rag_endpoints/endpoints.py
View file

@ -9,6 +9,7 @@ Provides:
import base64
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import orjson
@ -19,6 +20,9 @@ from starlette.datastructures import UploadFile
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import *
from litellm.proxy.auth.auth_utils import is_request_body_safe
@ -37,7 +41,7 @@ from litellm.proxy.rag_endpoints.upload_security import (
validate_upload,
)
from litellm.proxy.vector_store_endpoints.endpoints import (
_update_request_data_with_litellm_managed_vector_store_registry, # pyright: ignore[reportPrivateUsage] # shared registry-merge helper used by the direct search endpoint
build_request_data_from_managed_vector_store,
)
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
@ -123,12 +127,21 @@ def _collect_vector_store_ids_from_payload(payload: object) -> set[str]:
async def _authorize_nested_vector_store_ids(
payload: object,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)):
await assert_user_can_access_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
) -> Mapping[str, LiteLLM_ManagedVectorStore]:
"""Authorize every nested vector store id and return the managed stores it resolved."""
return MappingProxyType(
{
vector_store_id: store
for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload))
if (
store := await assert_user_can_access_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
)
is not None
}
)
def _build_file_metadata_entry(
@ -703,23 +716,24 @@ async def rag_query(
status_code=400,
detail={"error": "retrieval_config must contain 'vector_store_id'"},
)
await _authorize_nested_vector_store_ids(
resolved_stores: Final = await _authorize_nested_vector_store_ids(
payload=retrieval_config,
user_api_key_dict=user_api_key_dict,
)
# Merge litellm-managed vector store params (provider, region, embedding
# model, credentials, ...) from the registry — same source the direct
# /vector_stores/{id}/search endpoint uses. User-supplied
# retrieval_config keys win on conflict.
store_data: Final = await _update_request_data_with_litellm_managed_vector_store_registry(
data={}, # mutable-ok: the helper mutates and returns the seed dict
vector_store_id=retrieval_config["vector_store_id"],
user_api_key_dict=user_api_key_dict,
# model, credentials, ...) from the registry: the same source the direct
# /vector_stores/{id}/search endpoint uses. Store-managed keys win on
# conflict so callers cannot override the store's provider or credentials.
managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"])
store_data: Final = (
await build_request_data_from_managed_vector_store(managed_store)
if managed_store is not None
else MappingProxyType({})
)
merged_retrieval_config: Final = {
**store_data,
**retrieval_config,
**store_data,
} # mutable-ok: litellm.aquery requires a plain dict payload
# Add litellm data
@ -733,7 +747,12 @@ async def rag_query(
proxy_config=proxy_config,
)
verbose_proxy_logger.debug("RAG Query - model: %s, retrieval_config: %s", model, merged_retrieval_config)
verbose_proxy_logger.debug(
"RAG Query - model: %s, vector_store_id: %s, custom_llm_provider: %s",
model,
retrieval_config["vector_store_id"],
merged_retrieval_config.get("custom_llm_provider"),
)
# Call query
response: Final = await litellm.aquery(

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

@ -1,3 +1,5 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import (
Annotated,
Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
@ -32,6 +34,41 @@ router: Final = APIRouter()
########################################################
async def build_request_data_from_managed_vector_store(
vector_store: LiteLLM_ManagedVectorStore,
) -> Mapping[str, object]:
"""
Build request params (provider, credential ref, litellm_params) from an
already-resolved managed vector store.
``litellm_embedding_config`` is resolved here, at request-handling time,
instead of at row-creation time: the resolved api_key/api_base/api_version
lives only in the returned per-request mapping and is never persisted back
to the registry cache. Legacy rows that already carry a resolved
(cleartext) config skip the lookup and pass through unchanged.
"""
top_level: Final = MappingProxyType(
{
key: vector_store.get(key)
for key in ("custom_llm_provider", "litellm_credential_name")
if key in vector_store
}
)
litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({})
embedding_model: Final = litellm_params.get("litellm_embedding_model")
if not embedding_model or litellm_params.get("litellm_embedding_config"):
return MappingProxyType({**top_level, **litellm_params})
from litellm.proxy.proxy_server import prisma_client
resolved_config: Final = await _resolve_embedding_config(
embedding_model=embedding_model, prisma_client=prisma_client
)
if not resolved_config:
return MappingProxyType({**top_level, **litellm_params})
return MappingProxyType({**top_level, **litellm_params, "litellm_embedding_config": resolved_config})
async def _update_request_data_with_litellm_managed_vector_store_registry(
data: dict,
vector_store_id: str,
@ -51,47 +88,14 @@ async def _update_request_data_with_litellm_managed_vector_store_registry(
vector_store_to_run: Final[LiteLLM_ManagedVectorStore | None] = await get_litellm_managed_vector_store(
vector_store_id=vector_store_id
)
if vector_store_to_run is not None:
if user_api_key_dict is not None:
await assert_user_can_access_vector_store(
vector_store=vector_store_to_run,
user_api_key_dict=user_api_key_dict,
)
if "custom_llm_provider" in vector_store_to_run:
data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider")
if "litellm_credential_name" in vector_store_to_run:
data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name")
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: Final = 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: Final = 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
if vector_store_to_run is None:
return data
if user_api_key_dict is not None:
await assert_user_can_access_vector_store(
vector_store=vector_store_to_run,
user_api_key_dict=user_api_key_dict,
)
return {**data, **(await build_request_data_from_managed_vector_store(vector_store_to_run))}
@router.post(

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

@ -470,7 +470,7 @@ async def create_vector_store_in_db(
# 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``
# ``build_request_data_from_managed_vector_store``
# at request-handling time so the cleartext config exists only in
# per-request memory and never reaches the database.
if litellm_params:
@ -864,7 +864,7 @@ async def update_vector_store(
# 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``
# ``build_request_data_from_managed_vector_store``
# so this row only ever stores the user-supplied
# ``litellm_embedding_model`` reference.
if "litellm_params" in update_data:

25
litellm/rag/main.py
View file

@ -51,12 +51,19 @@ INGESTION_REGISTRY: Final[dict[str, type[BaseRAGIngestion]]] = {
"vertex_ai": VertexAIRAGIngestion,
}
# retrieval_config keys consumed by the query pipeline itself; everything else is
# forwarded to vector_stores.asearch as provider-specific params (e.g.
# aws_region_name, embedding_model, vector_bucket_name for S3 Vectors).
# `filters`/`retrieval_filter` are reserved for the explicit filter param.
_CONSUMED_RETRIEVAL_CONFIG_KEYS: Final = frozenset(
{"vector_store_id", "custom_llm_provider", "top_k", "filters", "retrieval_filter"}
# Only these retrieval_config keys are forwarded to vector_stores.asearch as
# provider-specific params. The explicit allowlist keeps caller-controlled
# connection overrides (api_base, api_key, ...) away from the search call,
# where they could redirect store credentials to an attacker-chosen host.
_FORWARDABLE_RETRIEVAL_CONFIG_KEYS: Final = frozenset(
{
"aws_region_name",
"vector_bucket_name",
"embedding_model",
"litellm_embedding_model",
"litellm_embedding_config",
"litellm_credential_name",
}
)
@ -233,10 +240,10 @@ async def _execute_query_pipeline(
raise ValueError("No query found in messages for RAG query")
# 2. Search vector store
# Forward provider-specific retrieval_config extras (region, embedding model,
# bucket, credentials refs, ...) to the search call; kwargs win on conflict.
# Forward allowlisted provider retrieval_config extras (region, embedding
# model, bucket, credential refs) to the search call; kwargs win on conflict.
provider_search_params: Final = MappingProxyType(
{k: v for k, v in retrieval_config.items() if k not in _CONSUMED_RETRIEVAL_CONFIG_KEYS}
{k: v for k, v in retrieval_config.items() if k in _FORWARDABLE_RETRIEVAL_CONFIG_KEYS}
)
forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs})
with _suppressed_sub_call_billing():

26
tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py
View file

@ -357,12 +357,9 @@ def test_rag_query_merges_managed_store_params(client_internal_user):
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new_callable=AsyncMock,
return_value=mock_response,
) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and stubs the access assert covered by auth tests
"litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id",
new=AsyncMock(),
), patch( # test-quality-ok: stubs the direct-endpoint access assert covered by auth tests
"litellm.proxy.vector_store_endpoints.endpoints.assert_user_can_access_vector_store",
new=AsyncMock(),
) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs
"litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store",
new=AsyncMock(return_value=True),
):
response = client_internal_user.post(
"/v1/rag/query",
@ -383,8 +380,8 @@ def test_rag_query_merges_managed_store_params(client_internal_user):
assert forwarded_config["vector_bucket_name"] == "bkt"
def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user):
"""User-supplied retrieval_config keys must win over registry values."""
def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_user):
"""Registry values must win over user-supplied retrieval_config keys so callers cannot override store credentials."""
import litellm
from litellm.types.utils import ModelResponse
@ -406,12 +403,9 @@ def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user):
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new_callable=AsyncMock,
return_value=mock_response,
) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and stubs the access assert covered by auth tests
"litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id",
new=AsyncMock(),
), patch( # test-quality-ok: stubs the direct-endpoint access assert covered by auth tests
"litellm.proxy.vector_store_endpoints.endpoints.assert_user_can_access_vector_store",
new=AsyncMock(),
) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs
"litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store",
new=AsyncMock(return_value=True),
):
response = client_internal_user.post(
"/v1/rag/query",
@ -424,7 +418,9 @@ def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user):
assert response.status_code == 200, response.json()
forwarded_config = mock_aquery.await_args.kwargs["retrieval_config"]
assert forwarded_config["aws_region_name"] == "us-east-1"
assert forwarded_config["aws_region_name"] == "eu-west-1"
EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*"
INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'

39
tests/test_litellm/rag/test_main.py
View file

@ -349,6 +349,45 @@ async def test_aquery_minimal_retrieval_config_forwards_no_extras():
assert not (leaked & set(search_kwargs.keys()))
@pytest.mark.asyncio
async def test_aquery_does_not_forward_connection_override_keys_to_search():
"""
Only allowlisted retrieval_config keys may reach the vector store search
call. Caller-controlled connection overrides (api_base, api_key, arbitrary
extras) must be dropped, otherwise a caller could redirect store
credentials to an attacker-chosen host.
"""
from unittest.mock import AsyncMock
from litellm.types.vector_stores import VectorStoreSearchResponse
fake_search = AsyncMock(
return_value=VectorStoreSearchResponse(
object="vector_store.search_results.page", search_query="q", data=[]
)
)
with patch("litellm.vector_stores.asearch", new=fake_search): # test-quality-ok: asearch is the boundary the forwarding contract under test targets
await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
retrieval_config={
"vector_store_id": "bkt:idx",
"custom_llm_provider": "s3_vectors",
"aws_region_name": "eu-west-1",
"api_base": "https://attacker.example.com",
"api_key": "attacker-key",
"arbitrary_extra": "nope",
},
mock_response="hi",
)
fake_search.assert_awaited_once()
search_kwargs = fake_search.await_args.kwargs
assert search_kwargs["aws_region_name"] == "eu-west-1"
blocked = {"api_base", "api_key", "arbitrary_extra"}
assert not (blocked & set(search_kwargs.keys()))
def test_rag_call_types_are_registered():
"""
query/aquery/ingest/aingest are @client-decorated entry points, so their