scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(rag): resolve vector_store_registry credentials in /v1/rag/query

Browse source
This commit is contained in:
dcsgod 2026-08-02 23:35:40 +05:30
parent 491eda319c
commit 5b02f731b3
4 changed files with 146 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -26,6 +26,9 @@ from litellm.proxy.common_utils.http_parsing_utils import (
_safe_get_request_headers,
get_form_data,
)
from litellm.proxy.vector_store_endpoints.endpoints import (
_update_request_data_with_litellm_managed_vector_store_registry,
)
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
)
@ -647,6 +650,11 @@ async def rag_query(
status_code=400,
detail={"error": "retrieval_config must contain 'vector_store_id'"},
)
retrieval_config = await _update_request_data_with_litellm_managed_vector_store_registry(
data=retrieval_config,
vector_store_id=retrieval_config["vector_store_id"],
user_api_key_dict=user_api_key_dict,
)
await _authorize_nested_vector_store_ids(
payload=retrieval_config,
user_api_key_dict=user_api_key_dict,

7
litellm/rag/main.py
View file

@ -226,13 +226,18 @@ async def _execute_query_pipeline(
raise ValueError("No query found in messages for RAG query")
# 2. Search vector store
_reserved_retrieval_keys = frozenset(
("vector_store_id", "top_k", "custom_llm_provider", "query", "max_num_results")
)
search_params = {k: v for k, v in {**retrieval_config, **kwargs}.items() if k not in _reserved_retrieval_keys}
with _suppressed_sub_call_billing():
search_response = await litellm.vector_stores.asearch(
vector_store_id=retrieval_config["vector_store_id"],
query=query_text,
max_num_results=retrieval_config.get("top_k", 10),
custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
**kwargs,
**search_params,
)
search_provider = retrieval_config.get("custom_llm_provider", "openai")

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

@ -13,9 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -62,19 +60,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
@ -101,8 +93,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()}"
)
@ -118,15 +109,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()}"
)
@ -174,13 +162,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
@ -197,9 +185,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": {
@ -219,9 +205,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):
@ -234,14 +218,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):
@ -265,12 +245,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",
@ -306,10 +288,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={
@ -327,3 +313,73 @@ def test_rag_query_stream_returns_event_stream(client_internal_user):
assert response.headers.get("content-type", "").startswith("text/event-stream")
assert '"object":"chat.completion.chunk"' in response.text
assert "data: [DONE]" in response.text
def test_rag_query_resolves_vector_store_registry_credentials(client_internal_user):
"""
/v1/rag/query must resolve vector store registry entries and pass the resolved
custom_llm_provider, api_key, api_base, etc. into litellm.vector_stores.asearch.
"""
import litellm
mock_vector_store = {
"vector_store_id": "vs_azure_123",
"custom_llm_provider": "azure_ai_search",
"litellm_credential_name": "azure_search_cred",
"litellm_params": {
"api_key": "fake-azure-search-key",
"api_base": "https://fake-azure-search.search.windows.net",
"api_version": "2023-11-01",
},
}
mock_registry = MagicMock()
mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store
mock_registry.pop_vector_stores_to_run_with_db_fallback = AsyncMock(return_value=None)
captured_asearch_kwargs = {}
async def fake_asearch(**kwargs):
captured_asearch_kwargs.update(kwargs)
mock_search_resp = MagicMock()
mock_search_resp.data = []
return mock_search_resp
mock_model_response = MagicMock(
id="chatcmpl-test",
choices=[
MagicMock(
message=MagicMock(role="assistant", content="Answer"),
finish_reason="stop",
)
],
_hidden_params={},
model="gpt-4o-mini",
)
with (
patch.object(litellm, "vector_store_registry", mock_registry),
patch(
"litellm.vector_stores.asearch",
new=AsyncMock(side_effect=fake_asearch),
),
patch(
"litellm.acompletion",
new=AsyncMock(return_value=mock_model_response),
),
):
response = client_internal_user.post(
"/v1/rag/query",
json={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "search test query"}],
"retrieval_config": {
"vector_store_id": "vs_azure_123",
},
},
)
assert response.status_code == 200, response.json()
assert captured_asearch_kwargs.get("custom_llm_provider") == "azure_ai_search"
assert captured_asearch_kwargs.get("api_key") == "fake-azure-search-key"
assert captured_asearch_kwargs.get("api_base") == "https://fake-azure-search.search.windows.net"
assert captured_asearch_kwargs.get("api_version") == "2023-11-01"

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

@ -11,7 +11,7 @@ aquery carries the completion response with real usage and cost.
"""
import asyncio
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import pytest
@ -89,6 +89,35 @@ async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use
assert standard_logging_object["response_cost"] > 0
@pytest.mark.asyncio
async def test_aquery_forwards_provider_credentials_to_vector_store_search():
"""Retrieval config credentials must be forwarded into the vector store search call."""
captured_kwargs = {}
async def fake_asearch(**kwargs):
captured_kwargs.update(kwargs)
return {"data": []}
with patch("litellm.vector_stores.asearch", new=AsyncMock(side_effect=fake_asearch)):
await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
retrieval_config={
"vector_store_id": "vs_test_123",
"custom_llm_provider": "azure_ai_search",
"api_key": "search-key",
"api_base": "https://example.com",
"api_version": "2024-01-01",
},
mock_response="hi there",
)
assert captured_kwargs["custom_llm_provider"] == "azure_ai_search"
assert captured_kwargs["api_key"] == "search-key"
assert captured_kwargs["api_base"] == "https://example.com"
assert captured_kwargs["api_version"] == "2024-01-01"
@pytest.mark.asyncio
async def test_aquery_response_hidden_params_carry_completion_cost():
"""