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 36

fix(rag): resolve vector store registry credentials in rag query

Browse source 
/v1/rag/query never resolved the LiteLLM-managed vector store registry,
so a vector store registered with a non-OpenAI provider (and its
litellm_params such as api_key and api_base) was still queried against
api.openai.com. The standalone /v1/vector_stores/{id}/search endpoint
already resolves this via _update_request_data_with_litellm_managed_
vector_store_registry; rag_query did not.

Reuse the shared resolver in rag_query to load the managed vector store
for the requested vector_store_id, set the resolved custom_llm_provider
on retrieval_config so _execute_query_pipeline picks it up, and forward
the resolved litellm_params (minus litellm_credential_name) through
request_data so they reach the vector store search.

Closes https://github.com/BerriAI/litellm/issues/35599
This commit is contained in:
Diwak4r 2026-08-23 00:07:39 +05:45
parent 28887f12c5
commit be116da8ee
2 changed files with 278 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -8,6 +8,7 @@ Provides:
import base64
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import orjson
@ -37,6 +38,43 @@ if TYPE_CHECKING:
router: Final = APIRouter()
# Fields that must never be logged verbatim from a client-supplied
# ``retrieval_config`` (provider-native stores legitimately carry connection
# credentials in the dict; the debug log below redacts them).
_SENSITIVE_RETRIEVAL_CONFIG_FIELDS: Final = frozenset(
{
"api_key",
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"vertex_credentials",
"vertex_ai_credentials",
"litellm_embedding_config",
}
)
# Connection credential/endpoint fields that an authenticated client must not
# be able to inject into the vector store search through ``retrieval_config``.
# For a managed store the operator-resolved registry params are authoritative
# (they travel on a dedicated channel); for a provider-native store connection
# params must come from the proxy's server-side configuration (env / litellm
# credentials), never from the request body - a client-supplied ``api_base``
# would let a tenant pivot the proxy's egress to an internal host (SSRF).
_RETRIEVAL_CONFIG_CLIENT_BLOCKED_FIELDS: Final = frozenset(
{
"api_key",
"api_base",
"api_version",
"aws_region_name",
"region_name",
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"litellm_embedding_config",
"litellm_embedding_model",
}
)
def _raise_vector_store_scan_depth_exceeded() -> None:
raise HTTPException(
@ -492,6 +530,7 @@ async def rag_ingest(
# Add litellm data
request_data: dict[str, Any] = {}
# rebind-ok: request_data is rebuilt from the litellm-enriched request.
request_data = await add_litellm_data_to_request(
data=request_data,
request=request,
@ -661,10 +700,17 @@ async def rag_query(
user_api_key_dict=user_api_key_dict,
)
# Deferred import: keeps vector_store_endpoints out of sys.modules at
# proxy boot so lazy router registration and the OpenAPI snapshot stub
# injection are not skipped for vector store routes.
from litellm.proxy.vector_store_endpoints.endpoints import (
_update_request_data_with_litellm_managed_vector_store_registry, # pyright: ignore[reportPrivateUsage] # one canonical managed-store registry resolver, shared with the vector store endpoints
)
# Add litellm data
request_data: dict[str, object] = {}
request_data = await add_litellm_data_to_request(
data=request_data,
data={}, # mutable-ok: resolver populates this dict in place
request=request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
@ -672,15 +718,51 @@ async def rag_query(
proxy_config=proxy_config,
)
verbose_proxy_logger.debug("RAG Query - model: %s, retrieval_config: %s", model, retrieval_config)
resolved_registry = await _update_request_data_with_litellm_managed_vector_store_registry(
data={}, # mutable-ok: resolver populates this dict in place
vector_store_id=retrieval_config["vector_store_id"],
user_api_key_dict=user_api_key_dict,
)
resolved_registry.pop("litellm_credential_name", None)
# An authenticated client must not control the search connection
# params: strip credential/endpoint fields from the client-supplied
# retrieval_config. Connection params reach the search only through
# the trusted managed-store registry channel below (or the proxy's
# own server-side configuration for provider-native stores).
search_retrieval_config: Final[Mapping[str, Any]] = MappingProxyType(
{
key: value
for key, value in retrieval_config.items()
if key not in _RETRIEVAL_CONFIG_CLIENT_BLOCKED_FIELDS
}
)
# Trusted registry params travel on a separate channel (never merged
# into the client-controlled ``retrieval_config``), so client-supplied
# credential/endpoint fields can never be forwarded to the search for
# a managed store, and resolved credentials are never logged. The
# search call site forwards only an allowlisted subset of this dict
# (see ``_VECTOR_STORE_SEARCH_PARAMS`` in ``litellm/rag/main.py``),
# so stored metadata / guardrail fields never reach the LLM call.
verbose_proxy_logger.debug(
"RAG Query - model: %s, retrieval_config: %s",
model,
MappingProxyType(
{
key: ("***" if key in _SENSITIVE_RETRIEVAL_CONFIG_FIELDS and value else value)
for key, value in search_retrieval_config.items()
}
),
)
# Call query
response: Final = await litellm.aquery(
model=model,
messages=messages,
retrieval_config=retrieval_config,
retrieval_config=search_retrieval_config,
rerank=rerank,
stream=stream,
litellm_managed_vector_store_registry=resolved_registry,
router=llm_router,
**request_data,
)

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

@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi.testclient import TestClient
import litellm
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import app
@ -57,19 +57,13 @@ def test_internal_user_viewer_rag_ingest_without_vector_store_id_rejected(
response = client_internal_user_viewer.post(
"/v1/rag/ingest",
files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")},
data={
"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'
},
data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'},
)
assert response.status_code == 403
detail = response.json()
assert "detail" in detail
error_msg = (
detail["detail"]["error"]
if isinstance(detail["detail"], dict)
else str(detail["detail"])
)
error_msg = detail["detail"]["error"] if isinstance(detail["detail"], dict) else str(detail["detail"])
assert "internal_user_viewer" in error_msg
assert "vector_store_id" in error_msg
@ -96,8 +90,7 @@ def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check(
# Should not be 403 (role check passed)
assert response.status_code != 403, (
f"internal_user_viewer with vector_store_id should pass role check. "
f"Response: {response.json()}"
f"internal_user_viewer with vector_store_id should pass role check. Response: {response.json()}"
)
@ -113,15 +106,12 @@ def test_internal_user_rag_ingest_without_vector_store_id_allowed(client_interna
response = client_internal_user.post(
"/v1/rag/ingest",
files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")},
data={
"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'
},
data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'},
)
# Should not be 403
assert response.status_code != 403, (
f"internal_user should be allowed to create new vector stores. "
f"Response: {response.json()}"
f"internal_user should be allowed to create new vector stores. Response: {response.json()}"
)
@ -169,13 +159,13 @@ def test_rag_ingest_blocks_clientside_credentials(client_internal_user, blocked_
},
},
)
assert (
response.status_code == 400
), f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}"
assert response.status_code == 400, (
f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}"
)
body = response.json()
assert blocked_field in str(
body
), f"Response should mention '{blocked_field}': {body}"
assert blocked_field in str(body), f"Response should mention '{blocked_field}': {body}"
class TestRagIngestSSRFBlocked:
"""
aws_sts_endpoint and related credential-redirect fields must be rejected
@ -192,9 +182,7 @@ class TestRagIngestSSRFBlocked:
("aws_bedrock_runtime_endpoint", "https://attacker.example/bedrock"),
],
)
def test_ssrf_field_in_vector_store_config_rejected(
self, field, value, client_internal_user
):
def test_ssrf_field_in_vector_store_config_rejected(self, field, value, client_internal_user):
payload = {
"file_url": "https://example.com/doc.pdf",
"ingest_options": {
@ -214,9 +202,7 @@ class TestRagIngestSSRFBlocked:
)
body = response.json()
detail = body.get("detail", {})
error_text = (
detail.get("error", "") if isinstance(detail, dict) else str(detail)
)
error_text = detail.get("error", "") if isinstance(detail, dict) else str(detail)
assert field in error_text, f"Error should name the offending field: {error_text}"
def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user):
@ -229,14 +215,10 @@ class TestRagIngestSSRFBlocked:
"/v1/rag/ingest",
json={
"file_url": "https://example.com/doc.pdf",
"ingest_options": {
"vector_store": {"custom_llm_provider": "bedrock"}
},
"ingest_options": {"vector_store": {"custom_llm_provider": "bedrock"}},
},
)
assert response.status_code != 400, (
f"Clean Bedrock ingest_options should not be rejected: {response.json()}"
)
assert response.status_code != 400, f"Clean Bedrock ingest_options should not be rejected: {response.json()}"
def test_rag_query_returns_response_cost_header(client_internal_user):
@ -260,12 +242,14 @@ def test_rag_query_returns_response_cost_header(client_internal_user):
)
mock_response._hidden_params["response_cost"] = 3.45e-06
with patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new_callable=AsyncMock,
return_value=mock_response,
), patch("litellm.vector_store_registry", None), patch(
"litellm.proxy.proxy_server.prisma_client", None
with (
patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new_callable=AsyncMock,
return_value=mock_response,
),
patch("litellm.vector_store_registry", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
):
response = client_internal_user.post(
"/v1/rag/query",
@ -283,6 +267,167 @@ def test_rag_query_returns_response_cost_header(client_internal_user):
assert response.headers.get("x-litellm-response-cost") == "3.45e-06"
@pytest.mark.asyncio
async def test_rag_query_resolves_managed_vector_store_registry(client_internal_user):
"""
/v1/rag/query must resolve the LiteLLM-managed vector store registry so the
configured provider and litellm_params (api_key, api_base) reach the
internal vector store search instead of defaulting to OpenAI.
Before the fix, retrieval_config["custom_llm_provider"] stayed whatever the
client sent and the registry credentials were never forwarded to
litellm.aquery, so a registered non-OpenAI vector store was always queried
via api.openai.com. The fix also keeps the registry params on a dedicated
channel (``litellm_managed_vector_store_registry``), separate from the
client-controlled ``retrieval_config``: stored fields such as metadata
must never surface as top-level kwargs, which flow to the generation call.
"""
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
from litellm.types.utils import ModelResponse
mock_vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "vs_registered",
"custom_llm_provider": "bedrock",
"litellm_credential_name": "my_bedrock_creds",
"litellm_params": {
"aws_region_name": "us-east-1",
# Security probe: a stored field that must never reach the
# generation-call kwargs / sanitization boundary.
"metadata": {"disable_global_guardrails": True},
},
}
mock_registry = MagicMock()
mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store
captured = {}
async def fake_aquery(**kwargs):
captured.update(kwargs)
return ModelResponse(
id="chatcmpl-test",
choices=[
{
"index": 0,
"message": {
"role": "assistant",
"content": "The codename is AZURE-FALCON-42.",
},
"finish_reason": "stop",
}
],
model="gpt-4o-mini",
usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49},
)
with (
patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new=AsyncMock(side_effect=fake_aquery),
),
patch.object(litellm, "vector_store_registry", mock_registry),
patch("litellm.proxy.proxy_server.prisma_client", None),
):
response = client_internal_user.post(
"/v1/rag/query",
json={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "What is the codename?"}],
"retrieval_config": {
"vector_store_id": "vs_registered",
"custom_llm_provider": "openai",
},
},
)
assert response.status_code == 200, response.text
# Registry params travel on the dedicated managed-store channel, separate
# from the client-supplied retrieval_config (which is left untouched)...
assert captured["retrieval_config"]["custom_llm_provider"] == "openai"
assert captured["retrieval_config"]["vector_store_id"] == "vs_registered"
assert "aws_region_name" not in captured["retrieval_config"]
resolved_registry = captured["litellm_managed_vector_store_registry"]
assert resolved_registry["custom_llm_provider"] == "bedrock"
assert resolved_registry["aws_region_name"] == "us-east-1"
assert resolved_registry["metadata"] == {"disable_global_guardrails": True}
# ...and never surface as top-level kwargs, which flow to the generation
# call. The credential name is popped (the search resolves it by id).
assert "aws_region_name" not in captured
# `metadata` at the top level is the request's own metadata (agent_id,
# endpoint, headers...); the registry's guardrail override must not leak
# into it, or the generation call would skip global guardrails.
assert "disable_global_guardrails" not in captured.get("metadata", {})
assert "litellm_credential_name" not in captured
assert "litellm_credential_name" not in captured["retrieval_config"]
assert "litellm_credential_name" not in resolved_registry
@pytest.mark.asyncio
async def test_rag_query_strips_client_connection_params(client_internal_user):
"""
Credential/endpoint fields (api_base, api_key, aws creds, ...) supplied by
the client inside retrieval_config must never reach the vector store
search: for a provider-native store they would let a tenant point the
proxy's egress at an arbitrary host (SSRF), and for a managed store the
operator-resolved registry params are authoritative. The proxy strips
these fields before calling litellm.aquery.
"""
from litellm.types.utils import ModelResponse
captured = {}
async def fake_aquery(**kwargs):
captured.update(kwargs)
return ModelResponse(
id="chatcmpl-test",
choices=[
{
"index": 0,
"message": {
"role": "assistant",
"content": "The codename is AZURE-FALCON-42.",
},
"finish_reason": "stop",
}
],
model="gpt-4o-mini",
usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49},
)
with (
patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new=AsyncMock(side_effect=fake_aquery),
),
patch("litellm.vector_store_registry", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
):
response = client_internal_user.post(
"/v1/rag/query",
json={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "What is the codename?"}],
"retrieval_config": {
"vector_store_id": "vs_test_123",
"custom_llm_provider": "openai",
"api_base": "http://attacker.example/internal",
"api_key": "sk-attacker",
"aws_secret_access_key": "AKIATTACKER",
},
},
)
assert response.status_code == 200, response.text
retrieval_config = captured["retrieval_config"]
assert retrieval_config["vector_store_id"] == "vs_test_123"
assert retrieval_config["custom_llm_provider"] == "openai"
assert "api_base" not in retrieval_config
assert "api_key" not in retrieval_config
assert "aws_secret_access_key" not in retrieval_config
def test_rag_query_stream_returns_event_stream(client_internal_user):
"""
A stream=true /v1/rag/query must return an SSE response. Returning the raw
@ -301,10 +446,14 @@ def test_rag_query_stream_returns_event_stream(client_internal_user):
api_key="test-key",
)
with patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new=AsyncMock(side_effect=fake_aquery),
), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None):
with (
patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new=AsyncMock(side_effect=fake_aquery),
),
patch("litellm.vector_store_registry", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
):
response = client_internal_user.post(
"/v1/rag/query",
json={