mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(rag): apply the managed vector store connection policy on the RAG routes and move the approval marker to constants
This commit is contained in:
mateo-berri
parent
db4e650f4a
commit
6fef6e47a5
5 changed files with 171 additions and 24 deletions
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 13427
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2187
|
||||
"limit": 2188
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 319
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44134
|
||||
"limit": 44131
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38255
|
||||
"limit": 38252
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19582
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29814
|
||||
"limit": 29806
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 110
|
||||
|
|
|
|||
|
|
@ -44,10 +44,12 @@ from litellm.proxy.rag_endpoints.upload_security import (
|
|||
validate_upload,
|
||||
)
|
||||
from litellm.proxy.vector_store_endpoints.endpoints import (
|
||||
apply_managed_vector_store_connection_policy,
|
||||
build_request_data_from_managed_vector_store,
|
||||
reject_caller_embedding_selection_params,
|
||||
)
|
||||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
assert_proxy_admin_for_user_supplied_vector_store_connection,
|
||||
assert_user_can_access_vector_store_id,
|
||||
)
|
||||
from litellm.repositories.table_repositories import ManagedVectorStoresRepository
|
||||
|
|
@ -134,7 +136,7 @@ async def _authorize_nested_vector_store_ids(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Mapping[str, LiteLLM_ManagedVectorStore]:
|
||||
"""Authorize every nested vector store id and return the managed stores it resolved."""
|
||||
return MappingProxyType(
|
||||
resolved_stores: Final = MappingProxyType(
|
||||
{
|
||||
vector_store_id: store
|
||||
for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload))
|
||||
|
|
@ -147,6 +149,14 @@ async def _authorize_nested_vector_store_ids(
|
|||
is not None
|
||||
}
|
||||
)
|
||||
for store in resolved_stores.values():
|
||||
assert_proxy_admin_for_user_supplied_vector_store_connection(
|
||||
custom_llm_provider=store.get("custom_llm_provider"),
|
||||
litellm_params=store.get("litellm_params"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
managed=True,
|
||||
)
|
||||
return resolved_stores
|
||||
|
||||
|
||||
def _build_file_metadata_entry(
|
||||
|
|
@ -727,20 +737,17 @@ async def rag_query(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Merge litellm-managed vector store params (provider, region, embedding
|
||||
# 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 = (
|
||||
build_request_data_from_managed_vector_store(managed_store)
|
||||
if managed_store is not None
|
||||
else MappingProxyType({})
|
||||
)
|
||||
merged_retrieval_config: Final = {
|
||||
**retrieval_config,
|
||||
**store_data,
|
||||
} # mutable-ok: litellm.aquery requires a plain dict payload
|
||||
merged_retrieval_config: Final = apply_managed_vector_store_connection_policy(
|
||||
data=retrieval_config,
|
||||
vector_store=managed_store,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Add litellm data
|
||||
request_data: dict[str, object] = {}
|
||||
|
|
|
|||
|
|
@ -96,7 +96,22 @@ 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 None:
|
||||
if vector_store_to_run is not None and 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 apply_managed_vector_store_connection_policy(
|
||||
data=data, vector_store=vector_store_to_run, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
|
||||
|
||||
def apply_managed_vector_store_connection_policy(
|
||||
data: Mapping[str, object],
|
||||
vector_store: LiteLLM_ManagedVectorStore | None,
|
||||
user_api_key_dict: UserAPIKeyAuth | None,
|
||||
) -> dict[str, object]:
|
||||
if vector_store is None:
|
||||
caller_data: Final = {key: value for key, value in data.items() if key != MILVUS_ADMIN_CONFIGURED_CONNECTION}
|
||||
if user_api_key_dict is not None:
|
||||
assert_proxy_admin_for_user_supplied_vector_store_connection(
|
||||
|
|
@ -105,18 +120,12 @@ async def _update_request_data_with_litellm_managed_vector_store_registry(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
return caller_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,
|
||||
)
|
||||
blocked_fields: Final = managed_connection_fields(
|
||||
vector_store_to_run.get("custom_llm_provider"), vector_store_to_run.get("litellm_params")
|
||||
vector_store.get("custom_llm_provider"), vector_store.get("litellm_params")
|
||||
)
|
||||
managed_data: Final = build_request_data_from_managed_vector_store(vector_store_to_run)
|
||||
request_data: Final = {
|
||||
**{key: value for key, value in data.items() if key not in blocked_fields},
|
||||
**managed_data,
|
||||
**build_request_data_from_managed_vector_store(vector_store),
|
||||
}
|
||||
return request_data
|
||||
|
||||
|
|
|
|||
|
|
@ -662,6 +662,137 @@ def test_rag_query_gates_managed_milvus_grpc_store_on_admin_approval(client_inte
|
|||
mock_aquery.assert_not_awaited()
|
||||
|
||||
|
||||
def _managed_grpc_store(vector_store_id: str, *, approved: bool) -> dict[str, object]:
|
||||
from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION
|
||||
|
||||
return {
|
||||
"vector_store_id": vector_store_id,
|
||||
"custom_llm_provider": "milvus",
|
||||
"litellm_params": {
|
||||
"custom_llm_provider": "milvus",
|
||||
"api_base": "http://127.0.0.1:19530",
|
||||
"api_key": "root:Milvus",
|
||||
"milvus_transport": "grpc",
|
||||
"milvus_text_field": "text",
|
||||
"litellm_embedding_model": "team-embedding",
|
||||
**({MILVUS_ADMIN_CONFIGURED_CONNECTION: True} if approved else {}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _registry_with(store: dict[str, object]) -> MagicMock:
|
||||
registry = MagicMock()
|
||||
registry.get_litellm_managed_vector_store_from_registry.return_value = store
|
||||
return registry
|
||||
|
||||
|
||||
def test_rag_query_rejects_a_non_admin_inline_grpc_connection_on_an_unregistered_store(client_internal_user):
|
||||
from litellm.llms.milvus.vector_stores.connection import MilvusConnectionRejection
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: aquery is the downstream boundary; the test asserts it is never reached
|
||||
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery", new_callable=AsyncMock
|
||||
) as mock_aquery,
|
||||
patch( # test-quality-ok: no registry so the id resolves as provider-native
|
||||
"litellm.vector_store_registry", None
|
||||
),
|
||||
patch( # test-quality-ok: proxy module global, no injection seam
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
),
|
||||
):
|
||||
response = client_internal_user.post(
|
||||
"/v1/rag/query",
|
||||
json={
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"retrieval_config": {
|
||||
"vector_store_id": "vs_unregistered",
|
||||
"custom_llm_provider": "milvus",
|
||||
"milvus_transport": "grpc",
|
||||
"api_base": "http://10.0.0.1:19530",
|
||||
"api_key": "root:Milvus",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert MilvusConnectionRejection.ADMIN_REQUIRED.value in response.text
|
||||
mock_aquery.assert_not_awaited()
|
||||
|
||||
|
||||
def test_rag_query_drops_caller_connection_fields_for_an_approved_managed_grpc_store(client_internal_user):
|
||||
import litellm
|
||||
from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
mock_response = ModelResponse(
|
||||
id="chatcmpl-test",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
with (
|
||||
patch( # test-quality-ok: aquery is the downstream boundary; the forwarded config is what the test asserts
|
||||
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery", new_callable=AsyncMock, return_value=mock_response
|
||||
) as mock_aquery,
|
||||
patch.object( # test-quality-ok: seeds the approved managed store the policy under test reads
|
||||
litellm, "vector_store_registry", _registry_with(_managed_grpc_store("proof_grpc", approved=True))
|
||||
),
|
||||
patch( # test-quality-ok: store access is not under test, so the request reaches the connection policy
|
||||
"litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", new=AsyncMock(return_value=True)
|
||||
),
|
||||
):
|
||||
response = client_internal_user.post(
|
||||
"/v1/rag/query",
|
||||
json={
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"retrieval_config": {
|
||||
"vector_store_id": "proof_grpc",
|
||||
"api_base": "http://10.0.0.1:19530",
|
||||
"api_key": "attacker",
|
||||
"milvus_db_name": "other_db",
|
||||
"top_k": 3,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
forwarded_config = mock_aquery.await_args.kwargs["retrieval_config"]
|
||||
assert forwarded_config["api_base"] == "http://127.0.0.1:19530"
|
||||
assert forwarded_config["api_key"] == "root:Milvus"
|
||||
assert "milvus_db_name" not in forwarded_config
|
||||
assert forwarded_config["top_k"] == 3
|
||||
assert forwarded_config[MILVUS_ADMIN_CONFIGURED_CONNECTION] is True
|
||||
|
||||
|
||||
def test_rag_ingest_rejects_an_unapproved_managed_grpc_store(client_internal_user):
|
||||
import litellm
|
||||
from litellm.llms.milvus.vector_stores.connection import MilvusConnectionRejection
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: aingest is the downstream boundary; the test asserts it is never reached
|
||||
"litellm.proxy.rag_endpoints.endpoints.litellm.aingest", new_callable=AsyncMock
|
||||
) as mock_aingest,
|
||||
patch.object( # test-quality-ok: seeds the unapproved managed store the policy under test reads
|
||||
litellm, "vector_store_registry", _registry_with(_managed_grpc_store("proof_legacy", approved=False))
|
||||
),
|
||||
patch( # test-quality-ok: store access is not under test, so the request reaches the connection policy
|
||||
"litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", new=AsyncMock(return_value=True)
|
||||
),
|
||||
):
|
||||
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":"milvus","vector_store_id":"proof_legacy"}}}'
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert MilvusConnectionRejection.ADMIN_SAVE_REQUIRED.value in response.text
|
||||
mock_aingest.assert_not_awaited()
|
||||
|
||||
|
||||
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"}}}'
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 22171
|
||||
"limit": 22172
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26728
|
||||
"limit": 26727
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 261
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue