From f0b056949f7ac92ac06c8dee91bc685161223763 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:51:00 +0000 Subject: [PATCH] fix(rag): keep forwarded provider keys on /rag/query and mixed JWT team claims working A forwarded x-api-key now counts as a caller credential, so the /rag/query search uses it over the team deployment as at the merge base. Sorting the JWT team claim compares entries as text, so a stray number or null no longer turns into a 401. A Gemini store with no proxy key is a 500 again, and a non-string vector_store_id keeps the merge base's 400 message. --- litellm/proxy/auth/handle_jwt.py | 2 +- litellm/proxy/rag_endpoints/endpoints.py | 28 +- litellm/rag/ingestion/gemini_ingestion.py | 6 +- .../test_rag_team_provider_credentials.py | 783 ++++++++++++++++-- tests/unit/proxy/auth/test_handle_jwt.py | 35 + .../proxy/rag_endpoints/test_rag_endpoints.py | 28 + .../unit/rag/ingestion/test_base_ingestion.py | 4 +- 7 files changed, 806 insertions(+), 80 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 18b20d24bd7..84dc1b4ac82 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1705,7 +1705,7 @@ class JWTAuthManager: return None, None any_claim_team_resolved = False - for team_id in sorted(team_ids): + for team_id in sorted(team_ids, key=str): try: team_object = await get_team_object( team_id=team_id, diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index b607cc3fa7f..9aea08f9e76 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -837,7 +837,7 @@ async def rag_query( if not isinstance(vector_store_id, str): raise HTTPException( status_code=400, - detail={"error": "retrieval_config must contain a string 'vector_store_id'"}, + detail={"error": "vector_store_id must be a non-empty string"}, ) reject_caller_embedding_selection_params(payload=retrieval_config, source="retrieval_config") resolved_stores: Final = await _authorize_nested_vector_store_ids( @@ -867,25 +867,25 @@ async def rag_query( {key: value for key, value in retrieval_config.items() if key not in _EXPLICIT_CREDENTIAL_KEYS} ) ) - team_credentials: Final = _team_provider_credentials( - effective_vector_store_config, - llm_router, - user_api_key_dict, - ) - vector_store_params: Final = ( - {**team_credentials, **managed_store_params} if managed_store is not None else team_credentials - ) - - # Add litellm data - request_data: dict[str, object] = {} - request_data = await add_litellm_data_to_request( - data=request_data, + request_data: Final[dict[str, object]] = await add_litellm_data_to_request( + data={}, request=request, general_settings=general_settings, user_api_key_dict=user_api_key_dict, version=version, proxy_config=proxy_config, ) + forwarded_credentials: Final = _without_blank_credentials( + {key: value for key, value in request_data.items() if key in _EXPLICIT_CREDENTIAL_KEYS} + ) + team_credentials: Final = _team_provider_credentials( + {**effective_vector_store_config, **forwarded_credentials}, + llm_router, + user_api_key_dict, + ) + vector_store_params: Final = ( + {**team_credentials, **managed_store_params} if managed_store is not None else team_credentials + ) verbose_proxy_logger.debug( "RAG Query - model: %s, vector_store_id: %s, custom_llm_provider: %s", diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 5c76e8a0d5a..4ee16685e03 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -18,7 +18,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.llms.gemini.common_utils import GeminiModelInfo -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: @@ -101,10 +101,10 @@ class GeminiRAGIngestion(BaseRAGIngestion): api_base: Final = cast(str | None, vector_store_config.get("api_base")) or GeminiModelInfo.get_api_base() if not api_key: - raise RAGIngestBadRequestError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search") + raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search") if not api_base: - raise RAGIngestBadRequestError("GEMINI_API_BASE is required") + raise ValueError("GEMINI_API_BASE is required") api_version: Final = "v1beta" base_url: Final = f"{api_base}/{api_version}" diff --git a/tests/integration/authorization/test_rag_team_provider_credentials.py b/tests/integration/authorization/test_rag_team_provider_credentials.py index 689d13fc929..d2c43ffc4b9 100644 --- a/tests/integration/authorization/test_rag_team_provider_credentials.py +++ b/tests/integration/authorization/test_rag_team_provider_credentials.py @@ -16,8 +16,10 @@ from types import MappingProxyType from typing import Final, Literal from urllib.parse import urlsplit +import anthropic import httpx import jwt +import openai import psutil import pytest import yaml @@ -60,12 +62,31 @@ BATCH_INPUT: Final = ("batch.jsonl", (json.dumps(dict(BATCH_INPUT_ROW)) + "\n"). FAILING_STORE_NAMES: Final = MappingProxyType({"rag-upstream-401": 401, "rag-upstream-429": 429}) OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -PREFIXED: Final = re.compile(r"^/(team|wildcard|shared|environment|credential|caller|gemini|azure)(/[^?]*)") +PREFIXED: Final = re.compile(r"^/(team|wildcard|shared|environment|credential|caller|gemini|azure|bedrock)(/[^?]*)") STORE_FILES: Final = re.compile(r"^/v1/vector_stores/([^/]+)/files$") FILE_CONTENT: Final = re.compile(r"^/v1/files/([^/]+)/content$") +FILE_ITEM: Final = re.compile(r"^/v1/files/([^/]+)$") +BATCH_ITEM: Final = re.compile(r"^/v1/batches/([^/]+)$") +BATCH_CANCEL: Final = re.compile(r"^/v1/batches/([^/]+)/cancel$") +BEDROCK_DATA_SOURCES: Final = re.compile(r"^/knowledgebases/([^/]+)/datasources/$") +BEDROCK_DATA_SOURCE: Final = re.compile(r"^/knowledgebases/([^/]+)/datasources/([^/]+)$") +EMPTY_KNOWLEDGE_BASE: Final = "EMPTYKB000" +BUCKETLESS_KNOWLEDGE_BASE: Final = "NOBUCKET00" +WEB_DATA_SOURCE: Final = "WEBSOURCE0" STORE_SEARCH: Final = re.compile(r"^(?:/v1|/openai)/vector_stores/([^/]+)/search$") GEMINI_UPLOAD: Final = re.compile(r"^/upload/v1beta/fileSearchStores/([^:]+):uploadToFileSearchStore$") +JWT_TEAM_COUNT: Final = 10 +AWS_AMBIENT_CREDENTIALS: Final = ( + "AWS_PROFILE", + "AWS_SESSION_TOKEN", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_ROLE_ARN", + "AWS_ROLE_NAME", + "AWS_SESSION_NAME", + "AWS_WEB_IDENTITY_TOKEN_FILE", +) + pytestmark: Final = pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS) Signature = tuple[str, str, str, str | None] @@ -100,8 +121,13 @@ class Tenants: wildcard_bare_key: str wildcard_prefixed_key: str wildcard_global_key: str + team_s: str team_s_key: str teamless_key: str + jwt_teams: tuple[str, ...] + jwt_team_model: str + jwt_multi_member: str + azure_restricted_key: str credential_name: str unreachable_credential_name: str @@ -115,10 +141,10 @@ class RagRig: signing_key: rsa.RSAPrivateKey scenario: Scenario - def jwt(self, subject: str, team_id: str) -> str: + def jwt(self, subject: str, *team_ids: str) -> str: issued_at: Final = int(time.time()) return jwt.encode( - {"sub": subject, "groups": [team_id], "iat": issued_at, "exp": issued_at + 300}, + {"sub": subject, "groups": list(team_ids), "iat": issued_at, "exp": issued_at + 300}, self.signing_key, algorithm="RS256", headers={"kid": JWT_KEY_ID}, @@ -181,44 +207,127 @@ def _vector_store_reply(request: Request, outage: threading.Event) -> Reply: ) -def _file_reply() -> Reply: - return _json_reply( - { - "id": f"file-{uuid.uuid4().hex}", - "object": "file", - "bytes": len(DOCUMENT[1]), - "created_at": int(time.time()), - "filename": DOCUMENT[0], - "purpose": "assistants", - "status": "processed", - } - ) +def _owned_by(prefix: str, identity: str) -> bool: + return identity.startswith((f"file-{prefix}-", f"batch_{prefix}_")) + + +def _file_object(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "file", + "bytes": len(DOCUMENT[1]), + "created_at": int(time.time()), + "filename": DOCUMENT[0], + "purpose": "assistants", + "status": "processed", + } + + +def _file_reply(prefix: str) -> Reply: + return _json_reply(_file_object(f"file-{prefix}-{uuid.uuid4().hex}")) def _batch_input_reply() -> Reply: return Reply(status=200, body=BATCH_INPUT[1]) -def _batch_reply(request: Request) -> Reply: +def _batch_object(identity: str, input_file_id: str, status: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": input_file_id, + "completion_window": "24h", + "status": status, + "created_at": int(time.time()), + "errors": None, + "output_file_id": None, + "error_file_id": None, + "request_counts": {"total": 0, "completed": 0, "failed": 0}, + "metadata": None, + } + + +def _batch_reply(request: Request, prefix: str) -> Reply: body: Final = _json_body(request) + return _json_reply( + _batch_object(f"batch_{prefix}_{uuid.uuid4().hex}", str(body.get("input_file_id")), "validating") + ) + + +def _list_reply() -> Reply: + return _json_reply({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}) + + +def _owned_reply(prefix: str, identity: str, payload: Mapping[str, JsonValue]) -> Reply: + return _json_reply(payload) if _owned_by(prefix, identity) else _error_reply(404) + + +def _openai_get(route: str, prefix: str) -> Reply: + if route in {"/v1/files", "/v1/batches"}: + return _list_reply() + content: Final = FILE_CONTENT.match(route) + if content is not None: + return _batch_input_reply() if _owned_by(prefix, content[1]) else _error_reply(404) + file_item: Final = FILE_ITEM.match(route) + if file_item is not None: + return _owned_reply(prefix, file_item[1], _file_object(file_item[1])) + batch_item: Final = BATCH_ITEM.match(route) + if batch_item is not None: + return _owned_reply(prefix, batch_item[1], _batch_object(batch_item[1], f"file-{prefix}-input", "validating")) + return _error_reply(404) + + +def _openai_delete(route: str, prefix: str) -> Reply: + file_item: Final = FILE_ITEM.match(route) + if file_item is None: + return _error_reply(404) + return _owned_reply(prefix, file_item[1], {"id": file_item[1], "object": "file", "deleted": True}) + + +def _embedding_reply(request: Request) -> Reply: + body: Final = _json_body(request) + inputs: Final = body.get("input") + count: Final = len(inputs) if isinstance(inputs, list) else 1 return _json_reply( { - "id": f"batch_{uuid.uuid4().hex}", - "object": "batch", - "endpoint": str(body.get("endpoint")), - "input_file_id": str(body.get("input_file_id")), - "completion_window": str(body.get("completion_window")), - "status": "validating", - "created_at": int(time.time()), - "errors": None, - "output_file_id": None, - "error_file_id": None, - "request_counts": {"total": 0, "completed": 0, "failed": 0}, - "metadata": None, + "object": "list", + "data": [{"object": "embedding", "index": index, "embedding": [0.1, 0.2, 0.3]} for index in range(count)], + "model": str(body.get("model")), + "usage": {"prompt_tokens": count, "total_tokens": count}, } ) +def _bedrock_reply(request: Request, route: str) -> Reply: + listed: Final = BEDROCK_DATA_SOURCES.match(route) + if request.method == "POST" and listed is not None: + summary: Final[dict[str, JsonValue]] = { + "knowledgeBaseId": listed[1], + "dataSourceId": WEB_DATA_SOURCE, + "name": "integration-web-source", + "status": "AVAILABLE", + "updatedAt": "2026-01-01T00:00:00Z", + } + return _json_reply({"dataSourceSummaries": [] if listed[1] == EMPTY_KNOWLEDGE_BASE else [summary]}) + detail: Final = BEDROCK_DATA_SOURCE.match(route) + if request.method == "GET" and detail is not None: + return _json_reply( + { + "dataSource": { + "knowledgeBaseId": detail[1], + "dataSourceId": detail[2], + "name": "integration-web-source", + "status": "AVAILABLE", + "dataSourceConfiguration": {"type": "WEB"}, + "createdAt": "2026-01-01T00:00:00Z", + "updatedAt": "2026-01-01T00:00:00Z", + } + } + ) + return _error_reply(404) + + def _attach_reply(request: Request, store_id: str) -> Reply: return _json_reply( { @@ -294,28 +403,59 @@ def _chat_reply(request: Request) -> Reply: def _responses_reply(request: Request) -> Reply: identity: Final = uuid.uuid4().hex - return _json_reply( - { - "id": f"resp_{identity}", - "object": "response", - "created_at": int(time.time()), - "status": "completed", - "model": str(_json_body(request).get("model")), - "output": [ - { - "type": "message", - "id": f"msg_{identity}", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "integration response", "annotations": []}], - } - ], - "parallel_tool_calls": False, - "tool_choice": "auto", - "tools": [], - "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, - } + completed: Final = _responses_object(identity, str(_json_body(request).get("model"))) + if _json_body(request).get("stream") is not True: + return _json_reply(completed) + events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( + ( + "response.created", + { + "type": "response.created", + "sequence_number": 0, + "response": {**completed, "status": "in_progress", "output": []}, + }, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": "integration response", + "logprobs": [], + }, + ), + ("response.completed", {"type": "response.completed", "sequence_number": 2, "response": completed}), ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events), + ) + + +def _responses_object(identity: str, model: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{identity}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": model, + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "integration response", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } def _gemini_reply(request: Request, route: str) -> Reply: @@ -330,21 +470,28 @@ def _gemini_reply(request: Request, route: str) -> Reply: return _error_reply(404) -def _openai_reply(request: Request, route: str, outage: threading.Event) -> Reply: +def _openai_reply(request: Request, prefix: str, route: str, outage: threading.Event) -> Reply: if request.method == "GET": - return _batch_input_reply() if FILE_CONTENT.match(route) is not None else _error_reply(404) + return _openai_get(route, prefix) + if request.method == "DELETE": + return _openai_delete(route, prefix) if request.method != "POST": return _error_reply(404) if route == "/v1/vector_stores": return _vector_store_reply(request, outage) - if route == "/v1/files": - return _file_reply() + if route in {"/v1/files", "/openai/files"}: + return _file_reply(prefix) if route == "/v1/batches": - return _batch_reply(request) + return _batch_reply(request, prefix) + cancelled: Final = BATCH_CANCEL.match(route) + if cancelled is not None: + return _owned_reply(prefix, cancelled[1], _batch_object(cancelled[1], f"file-{prefix}-input", "cancelling")) if route == "/v1/chat/completions": return _chat_reply(request) if route == "/v1/responses": return _responses_reply(request) + if route == "/v1/embeddings": + return _embedding_reply(request) attach: Final = STORE_FILES.match(route) if attach is not None: return _attach_reply(request, attach[1]) @@ -360,7 +507,9 @@ def _responder(outage: threading.Event) -> Callable[[Request], Reply]: return _error_reply(404) if matched[1] == "gemini": return _gemini_reply(request, matched[2]) - return _openai_reply(request, matched[2], outage) + if matched[1] == "bedrock": + return _bedrock_reply(request, matched[2]) + return _openai_reply(request, matched[1], matched[2], outage) return respond @@ -375,7 +524,7 @@ def _proxy_config(directory: Path) -> Path: "enable_jwt_auth": True, "litellm_jwtauth": {"user_id_jwt_field": "sub", "team_ids_jwt_field": "groups"}, }, - "files_settings": [{"custom_llm_provider": "openai"}], + "files_settings": [{"custom_llm_provider": "openai"}, {"custom_llm_provider": "azure"}], } path: Final = directory / "rag_team_provider_credentials.yaml" path.write_text(yaml.safe_dump(config)) @@ -397,6 +546,38 @@ def _provider_environment(upstream: str) -> Mapping[str, str]: ) +def _aws_environment(upstream: str) -> Mapping[str, str]: + return MappingProxyType( + { + "AWS_ENDPOINT_URL_BEDROCK_AGENT": f"{upstream}/bedrock", + "AWS_ACCESS_KEY_ID": "integration-rag-aws-access-key", + "AWS_SECRET_ACCESS_KEY": "integration-rag-aws-secret-key", + "AWS_REGION_NAME": "us-east-1", + "AWS_DEFAULT_REGION": "us-east-1", + "AWS_EC2_METADATA_DISABLED": "true", + } + ) + + +def _jwt_team_key(team_id: str) -> str: + return f"integration-rag-jwt-{team_id}" + + +def _jwt_team(gateway: Gateway, scenario: Scenario, upstream: str, model: str, member: str) -> str: + team_id: Final = scenario.team(models=[model]) + _deployment( + gateway, + scenario, + name=model, + model="openai/gpt-4o-mini", + api_key=_jwt_team_key(team_id), + api_base=f"{upstream}/team/v1", + team_id=team_id, + ) + gateway.post("/team/member_add", {"team_id": team_id, "member": {"user_id": member, "role": "user"}}) + return team_id + + def _deployment( gateway: Gateway, scenario: Scenario, @@ -486,6 +667,18 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants: team_id=team_w, ) team_s: Final = scenario.team(models=[shared_model]) + team_z: Final = scenario.team(models=[shared_model]) + _deployment( + gateway, + scenario, + name="azure/*", + model="azure/*", + api_key=WILDCARD_API_KEY, + api_base=f"{upstream}/azure", + team_id=team_z, + ) + jwt_team_model: Final = f"integration-rag-jwt-{uuid.uuid4().hex}" + jwt_multi_member: Final = scenario.user(user_role="internal_user") return Tenants( shared_model=shared_model, global_model=global_model, @@ -500,8 +693,15 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants: wildcard_bare_key=scenario.key(team_id=team_w, models=[BARE_WILDCARD_MODEL]), wildcard_prefixed_key=scenario.key(team_id=team_w, models=[WILDCARD_MODEL]), wildcard_global_key=scenario.key(team_id=team_w, models=[global_model]), + team_s=team_s, team_s_key=scenario.key(team_id=team_s), teamless_key=scenario.key(), + jwt_teams=tuple( + _jwt_team(gateway, scenario, upstream, jwt_team_model, jwt_multi_member) for _ in range(JWT_TEAM_COUNT) + ), + jwt_team_model=jwt_team_model, + jwt_multi_member=jwt_multi_member, + azure_restricted_key=scenario.key(team_id=team_z, models=["azure/gpt-4o-mini"]), credential_name=_credential(gateway, scenario, f"{upstream}/credential/v1"), unreachable_credential_name=_credential(gateway, scenario, _closed_url()), ) @@ -540,9 +740,13 @@ def rag_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RagRig]: owned_proxy( admin_gateway, directory, - {**_provider_environment(upstream.url), "JWT_PUBLIC_KEY_URL": jwks.url}, + { + **_provider_environment(upstream.url), + **_aws_environment(upstream.url), + "JWT_PUBLIC_KEY_URL": jwks.url, + }, config=_proxy_config(directory), - remove_environment=("GOOGLE_API_KEY",), + remove_environment=("GOOGLE_API_KEY", *AWS_AMBIENT_CREDENTIALS), workers=2, ) as gateway, gateway.scenario() as owned_scenario, @@ -550,6 +754,12 @@ def rag_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RagRig]: yield RagRig(gateway, upstream, outage, tenants, signing_key, owned_scenario) +@pytest.fixture(autouse=True) +def _drain_calls_an_earlier_failed_cell_left(request: pytest.FixtureRequest) -> None: + if "rag_rig" in request.fixturenames: + request.getfixturevalue("rag_rig").upstream.drain() + + def _ingest_request(name: str, vector_store: Mapping[str, JsonValue]) -> Mapping[str, str]: return MappingProxyType( {"request": json.dumps({"ingest_options": {"name": name, "vector_store": dict(vector_store)}})} @@ -1056,7 +1266,7 @@ def test_non_string_vector_store_id_is_rejected_before_search(rag_rig: RagRig, v ) assert response.status_code == 400, response.text - assert "retrieval_config must contain a string 'vector_store_id'" in response.text, response.text + assert "vector_store_id must be a non-empty string" in response.text, response.text assert rag_rig.calls() == () @@ -1234,7 +1444,7 @@ async def test_worker_sigkill_mid_query_burst_leaves_the_sibling_on_team_credent held_searches.put(request.target) assert release.wait(timeout=60), "The burst was never released" return _search_reply(request) - return _openai_reply(request, matched[2], threading.Event()) + return _openai_reply(request, matched[1], matched[2], threading.Event()) with ( gateway_from_environment() as admin_gateway, @@ -1364,3 +1574,456 @@ def test_a_team_with_two_openai_deployments_keeps_ingesting_through_the_older_on assert [calls for _, calls in ingests] == [ _ingest_signatures(picked_prefix, picked_key, store_id) for store_id, _ in ingests ] + + +@pytest.mark.parametrize( + ("stored", "prefix", "api_key"), + (("own_key", "caller", CALLER_API_KEY), ("null_key", "team", TEAM_API_KEY)), + ids=("own_key", "null_key"), +) +def test_ingest_into_a_managed_store_uses_its_own_key_else_the_team_key( + rag_rig: RagRig, stored: Literal["own_key", "null_key"], prefix: str, api_key: str +) -> None: + litellm_params: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "own_key": {"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"}, + "null_key": {"api_key": None}, + }[stored] + ) + with rag_rig.gateway.scenario() as scenario: + store_id: Final = _managed_store(rag_rig, scenario, litellm_params) + + response: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, "vector_store_id": store_id}) + + assert response.status_code == 200, response.text + assert object_value(response.json())["vector_store_id"] == store_id, response.text + bearer: Final = f"Bearer {api_key}" + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + (prefix, "POST", f"/v1/vector_stores/{store_id}/files", bearer), + ) + + +def _file_id_ingest(gateway: Gateway, key: str, vector_store: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", + "/v1/rag/ingest", + { + "file_id": "file-team-existing", + "ingest_options": {"name": "litellm-rag-ingest", "vector_store": dict(vector_store)}, + }, + key=key, + ) + + +@pytest.mark.parametrize( + ("vector_store", "ingestion"), + ( + ({"custom_llm_provider": "gemini"}, "GeminiRAGIngestion"), + ({"custom_llm_provider": "bedrock", "vector_store_id": BUCKETLESS_KNOWLEDGE_BASE}, "BedrockRAGIngestion"), + ({"custom_llm_provider": "s3_vectors", "vector_bucket_name": "integration-bucket"}, "S3VectorsRAGIngestion"), + ({"custom_llm_provider": "vertex_ai"}, "VertexAIRAGIngestion"), + ), + ids=("gemini", "bedrock", "s3_vectors", "vertex_ai"), +) +def test_existing_file_id_on_a_store_that_cannot_take_one_is_rejected_before_any_upstream_call( + rag_rig: RagRig, vector_store: Mapping[str, JsonValue], ingestion: str +) -> None: + before: Final = _store_rows("") + + response: Final = _file_id_ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, vector_store) + + assert response.status_code == 400, response.text + assert f"{ingestion} does not support ingesting an existing file_id" in response.text, response.text + assert rag_rig.calls() == () + assert _store_rows("") == before + + +def test_openai_file_id_without_a_vector_store_id_is_rejected_before_any_upstream_call(rag_rig: RagRig) -> None: + before: Final = _store_rows("") + + response: Final = _file_id_ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, OPENAI) + + assert response.status_code == 400, response.text + assert "vector_store_id is required when ingesting an existing file_id" in response.text, response.text + assert rag_rig.calls() == () + assert _store_rows("") == before + + +def test_openai_file_id_attaches_to_an_existing_store_with_the_team_key(rag_rig: RagRig) -> None: + store_id: Final = f"vs_existing_{uuid.uuid4().hex}" + + response: Final = _file_id_ingest( + rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, "vector_store_id": store_id} + ) + + assert _completed_store(rag_rig, response) == store_id + assert rag_rig.signatures() == (("team", "POST", f"/v1/vector_stores/{store_id}/files", f"Bearer {TEAM_API_KEY}"),) + + +@pytest.mark.parametrize( + ("knowledge_base", "error", "routes"), + ( + ( + EMPTY_KNOWLEDGE_BASE, + f"No data sources found for Knowledge Base {EMPTY_KNOWLEDGE_BASE}", + (("POST", f"/knowledgebases/{EMPTY_KNOWLEDGE_BASE}/datasources/"),), + ), + ( + BUCKETLESS_KNOWLEDGE_BASE, + f"Could not auto-detect S3 bucket for data source {WEB_DATA_SOURCE}", + ( + ("POST", f"/knowledgebases/{BUCKETLESS_KNOWLEDGE_BASE}/datasources/"), + ("GET", f"/knowledgebases/{BUCKETLESS_KNOWLEDGE_BASE}/datasources/{WEB_DATA_SOURCE}"), + ), + ), + ), + ids=("no_data_source", "no_s3_bucket"), +) +def test_bedrock_knowledge_base_the_caller_named_without_a_usable_data_source_fails_ingest_with_400( + rag_rig: RagRig, knowledge_base: str, error: str, routes: tuple[tuple[str, str], ...] +) -> None: + before: Final = _store_rows("") + + response: Final = _ingest( + rag_rig.gateway, + rag_rig.tenants.team_a_key, + {"custom_llm_provider": "bedrock", "vector_store_id": knowledge_base}, + ) + + assert response.status_code == 400, response.text + assert error in response.text, response.text + assert tuple((call.prefix, call.method, call.route) for call in rag_rig.calls()) == tuple( + ("bedrock", method, route) for method, route in routes + ) + assert _store_rows("") == before + + +def test_s3_vectors_bucket_name_the_caller_sent_is_rejected_with_400_after_embedding(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + before: Final = _store_rows("") + request: Final = { + "request": json.dumps( + { + "ingest_options": { + "name": "litellm-rag-ingest", + "embedding": {"model": tenants.shared_model}, + "vector_store": {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "ab", "dimension": 3}, + } + } + ) + } + + response: Final = rag_rig.gateway.request_multipart( + "/v1/rag/ingest", request, {"file": DOCUMENT}, key=tenants.team_a_key + ) + + assert response.status_code == 400, response.text + assert "Invalid vector_bucket_name 'ab'" in response.text, response.text + assert rag_rig.signatures() == (("shared", "POST", "/v1/embeddings", f"Bearer {SHARED_API_KEY}"),) + assert _store_rows("") == before + + +@dataclass(frozen=True, slots=True) +class KeylessRig: + gateway: Gateway + upstream: Wire + key: str + + +@pytest.fixture(scope="module") +def keyless_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[KeylessRig]: + directory: Final = tmp_path_factory.mktemp("rag_keyless_providers") + with ( + gateway_from_environment() as admin_gateway, + admin_gateway.scenario() as scenario, + wire_server(_responder(threading.Event())) as upstream, + ): + key: Final = scenario.key() + with owned_proxy( + admin_gateway, + directory, + {"OPENAI_API_KEY": ENVIRONMENT_API_KEY, "OPENAI_API_BASE": f"{upstream.url}/environment/v1"}, + remove_environment=("GEMINI_API_KEY", "GOOGLE_API_KEY", "VERTEXAI_PROJECT"), + workers=2, + ) as gateway: + yield KeylessRig(gateway, upstream, key) + + +@pytest.mark.parametrize( + ("provider", "status", "error"), + ( + ("gemini", 500, "GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search"), + ("vertex_ai", 400, "vertex_project is required for Vertex AI RAG ingestion"), + ), + ids=("gemini_without_a_proxy_key", "vertex_ai_without_a_project"), +) +def test_ingest_status_says_whether_the_caller_or_the_proxy_must_fix_the_missing_setting( + keyless_rig: KeylessRig, provider: str, status: int, error: str +) -> None: + before: Final = _store_rows("") + + response: Final = _ingest(keyless_rig.gateway, keyless_rig.key, {"custom_llm_provider": provider}) + + assert response.status_code == status, response.text + assert error in response.text, response.text + assert tuple(keyless_rig.upstream.drain()) == () + assert _store_rows("") == before + + +async def _jwt_chat_burst(rig: RagRig, token: str, model: str, count: int) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=str(rig.gateway.client.base_url), timeout=60, trust_env=False) as client: + return tuple( + await asyncio.gather( + *( + client.post( + "/v1/chat/completions", + json=_llm_body("/v1/chat/completions", model), + headers={"Authorization": f"Bearer {token}"}, + ) + for _ in range(count) + ) + ) + ) + + +async def test_jwt_in_ten_teams_that_hold_the_model_always_resolves_the_lowest_team_id(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + token: Final = rag_rig.jwt(tenants.jwt_multi_member, *reversed(tenants.jwt_teams)) + + responses: Final = await _jwt_chat_burst(rag_rig, token, tenants.jwt_team_model, 16) + + assert [response.status_code for response in responses] == [200] * 16, [response.text for response in responses] + assert ( + rag_rig.signatures() + == (("team", "POST", "/v1/chat/completions", f"Bearer {_jwt_team_key(min(tenants.jwt_teams))}"),) * 16 + ) + + +def _prompt() -> str: + return f"hello {uuid.uuid4().hex}" + + +def _sync_sdk_text(route: str, base_url: str, token: str, model: str, stream: bool) -> str: + if route == "/v1/messages": + messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=token, max_retries=0) + request: Final = {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": _prompt()}]} + if stream: + return "".join( + event.delta.text + for event in messages_client.messages.create(**request, stream=True) + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ) + message: Final = messages_client.messages.create(**request) + return "".join(block.text for block in message.content if block.type == "text") + client: Final = openai.OpenAI(base_url=f"{base_url}/v1", api_key=token, max_retries=0) + if route == "/v1/responses": + if stream: + return "".join( + event.delta + for event in client.responses.create(model=model, input=_prompt(), stream=True) + if event.type == "response.output_text.delta" + ) + return client.responses.create(model=model, input=_prompt()).output_text + chat_messages: Final = [{"role": "user", "content": _prompt()}] + if stream: + return "".join( + chunk.choices[0].delta.content or "" + for chunk in client.chat.completions.create(model=model, messages=chat_messages, stream=True) + if chunk.choices + ) + return client.chat.completions.create(model=model, messages=chat_messages).choices[0].message.content or "" + + +async def _async_sdk_text(route: str, base_url: str, token: str, model: str, stream: bool) -> str: + if route == "/v1/messages": + messages_client: Final = anthropic.AsyncAnthropic(base_url=base_url, api_key=token, max_retries=0) + request: Final = {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": _prompt()}]} + if stream: + return "".join( + [ + event.delta.text + async for event in await messages_client.messages.create(**request, stream=True) + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ] + ) + message: Final = await messages_client.messages.create(**request) + return "".join(block.text for block in message.content if block.type == "text") + client: Final = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=token, max_retries=0) + if route == "/v1/responses": + if stream: + return "".join( + [ + event.delta + async for event in await client.responses.create(model=model, input=_prompt(), stream=True) + if event.type == "response.output_text.delta" + ] + ) + return (await client.responses.create(model=model, input=_prompt())).output_text + chat_messages: Final = [{"role": "user", "content": _prompt()}] + if stream: + return "".join( + [ + chunk.choices[0].delta.content or "" + async for chunk in await client.chat.completions.create( + model=model, messages=chat_messages, stream=True + ) + if chunk.choices + ] + ) + completion: Final = await client.chat.completions.create(model=model, messages=chat_messages) + return completion.choices[0].message.content or "" + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +@pytest.mark.parametrize("client", ("sync", "async")) +@pytest.mark.parametrize("route", ("/v1/chat/completions", "/v1/messages", "/v1/responses")) +async def test_jwt_team_member_reaches_its_llm_routes_through_the_sdks( + rag_rig: RagRig, route: str, client: Literal["sync", "async"], stream: bool +) -> None: + tenants: Final = rag_rig.tenants + base_url: Final = str(rag_rig.gateway.client.base_url).rstrip("/") + token: Final = rag_rig.jwt(tenants.jwt_member, tenants.team_a) + + text: Final = ( + await asyncio.to_thread(_sync_sdk_text, route, base_url, token, tenants.team_a_model, stream) + if client == "sync" + else await _async_sdk_text(route, base_url, token, tenants.team_a_model, stream) + ) + + assert text == "integration response", text + calls: Final = rag_rig.calls() + assert [(call.prefix, call.authorization) for call in calls] == [("team", f"Bearer {TEAM_API_KEY}")], calls + + +def _batch_for(rig: RagRig, key: str) -> tuple[str, str]: + uploaded: Final = rig.gateway.request_multipart("/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = string_value(object_value(uploaded.json())["id"]) + created: Final = rig.gateway.request( + "POST", + "/v1/batches", + {"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + key=key, + ) + assert created.status_code == 200, created.text + rig.upstream.drain() + return file_id, string_value(object_value(created.json())["id"]) + + +@pytest.mark.parametrize( + ("method", "path", "route"), + ( + ("GET", "/v1/files/{file_id}", "/v1/files/{file_id}"), + ("GET", "/v1/files/{file_id}/content", "/v1/files/{file_id}/content"), + ("DELETE", "/v1/files/{file_id}", "/v1/files/{file_id}"), + ("GET", "/v1/files?purpose=batch&custom_llm_provider=openai", "/v1/files"), + ("GET", "/v1/batches/{batch_id}", "/v1/batches/{batch_id}"), + ("POST", "/v1/batches/{batch_id}/cancel", "/v1/batches/{batch_id}/cancel"), + ), + ids=("file_retrieve", "file_content", "file_delete", "file_list", "batch_retrieve", "batch_cancel"), +) +def test_files_and_batches_calls_by_a_key_restricted_under_the_team_wildcard_use_the_team_deployment( + rag_rig: RagRig, method: str, path: str, route: str +) -> None: + key: Final = rag_rig.tenants.wildcard_bare_key + file_id, batch_id = _batch_for(rag_rig, key) + + response: Final = rag_rig.gateway.request(method, path.format(file_id=file_id, batch_id=batch_id), key=key) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + ("wildcard", method, route.format(file_id=file_id, batch_id=batch_id), f"Bearer {WILDCARD_API_KEY}"), + ) + + +def test_batch_list_by_a_key_restricted_under_the_team_wildcard_is_answered_by_the_proxy_without_an_upstream_call( + rag_rig: RagRig, +) -> None: + key: Final = rag_rig.tenants.wildcard_bare_key + _batch_for(rag_rig, key) + + response: Final = rag_rig.gateway.request("GET", "/v1/batches?custom_llm_provider=openai", key=key) + + assert response.status_code == 200, response.text + assert object_value(response.json())["object"] == "list", response.text + assert rag_rig.signatures() == () + + +def _upload_batch_input(rig: RagRig, key: str) -> str: + uploaded: Final = rig.gateway.request_multipart("/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key) + assert uploaded.status_code == 200, uploaded.text + return string_value(object_value(uploaded.json())["id"]) + + +def _create_batch(rig: RagRig, key: str, file_id: str) -> httpx.Response: + return rig.gateway.request( + "POST", + "/v1/batches", + {"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + key=key, + ) + + +@pytest.mark.parametrize( + ("team", "prefix", "api_key"), + (("team_a", "team", TEAM_API_KEY), ("team_s", "shared", SHARED_API_KEY)), + ids=("team_deployment", "accessible_deployment"), +) +def test_batch_over_the_key_tpm_limit_is_rejected_after_reading_its_input_where_it_was_uploaded( + rag_rig: RagRig, team: Literal["team_a", "team_s"], prefix: str, api_key: str +) -> None: + key: Final = rag_rig.scenario.key(team_id=getattr(rag_rig.tenants, team), tpm_limit=1) + file_id: Final = _upload_batch_input(rag_rig, key) + + response: Final = _create_batch(rag_rig, key, file_id) + + assert response.status_code == 429, response.text + bearer: Final = f"Bearer {api_key}" + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + (prefix, "GET", f"/v1/files/{file_id}/content", bearer), + ) + + +def test_batch_whose_input_names_a_model_outside_the_key_allowlist_is_rejected_before_creation( + rag_rig: RagRig, +) -> None: + key: Final = rag_rig.tenants.restricted_key + file_id: Final = _upload_batch_input(rag_rig, key) + + response: Final = _create_batch(rag_rig, key, file_id) + + assert response.status_code == 403, response.text + assert BARE_WILDCARD_MODEL in response.text, response.text + bearer: Final = f"Bearer {SHARED_API_KEY}" + assert rag_rig.signatures() == ( + ("shared", "POST", "/v1/files", bearer), + ("shared", "GET", f"/v1/files/{file_id}/content", bearer), + ) + + +def test_azure_files_upload_by_a_key_restricted_under_the_team_azure_wildcard_uses_the_team_deployment( + rag_rig: RagRig, +) -> None: + response: Final = rag_rig.gateway.request_multipart( + "/v1/files", + {"purpose": "assistants", "custom_llm_provider": "azure"}, + {"file": DOCUMENT}, + key=rag_rig.tenants.azure_restricted_key, + ) + + assert response.status_code == 200, response.text + assert tuple((call.prefix, call.method, call.route, call.azure_key) for call in rag_rig.calls()) == ( + ("azure", "POST", "/openai/files", WILDCARD_API_KEY), + ) + + +def test_query_without_a_vector_store_id_is_rejected_before_search(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + response: Final = _query(rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, OPENAI) + + assert response.status_code == 400, response.text + assert "vector_store_id" in response.text, response.text + assert rag_rig.calls() == () diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 2b9013a4efa..d79c31d779f 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -338,6 +338,41 @@ async def test_a_jwt_in_two_teams_resolves_the_same_team_in_any_claim_order(clai assert team_id == "team-a" assert team_obj == teams["team-a"] + +@pytest.mark.asyncio +@pytest.mark.parametrize("stray_claim", [1, None]) +async def test_a_jwt_team_claim_with_a_non_string_entry_still_resolves_the_string_team(stray_claim: object): + jwt_handler: Final = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") + team: Final = LiteLLM_TeamTable(team_id="aa-team", models=["gpt-4"], metadata={}) + + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + side_effect=lambda team_id, **_: {"aa-team": team}[team_id], + ), + patch( + "litellm.proxy.auth.handle_jwt.can_team_access_model", + new_callable=AsyncMock, + return_value=True, + ), + ): + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids=set(jwt_handler.get_team_ids_from_jwt({"groups": ["aa-team", stray_claim]})), + requested_model="gpt-4", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert team_id == "aa-team" + assert team_obj == team + + _AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = { "test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": { "endpoint_id": "test-uuid-1", diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 6a4525560f9..d8247764a16 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -385,6 +385,7 @@ def _post_rag_query( retrieval_config: Mapping[str, object], managed_stores: Sequence[LiteLLM_ManagedVectorStore], router: Router, + headers: Mapping[str, str] | None = None, ) -> httpx.Response: with ( patch.object( # test-quality-ok: seeds the real managed-store registry the query path reads @@ -402,6 +403,7 @@ def _post_rag_query( "messages": [{"role": "user", "content": "hello"}], "retrieval_config": dict(retrieval_config), }, + headers=dict(headers or {}), ) @@ -688,6 +690,32 @@ def test_rag_query_without_a_team_deployment_keeps_the_environment_openai_creden assert search.calls.last.request.headers["authorization"] == f"Bearer {_ENVIRONMENT_OPENAI_KEY}" +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_search_uses_the_forwarded_provider_key_over_team_provider_credentials( + client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + forwarded_provider_key: Final = "sk-forwarded-provider-key" + monkeypatch.setenv("OPENAI_API_KEY", _ENVIRONMENT_OPENAI_KEY) + for unset in ("OPENAI_BASE_URL", "OPENAI_API_BASE"): + monkeypatch.delenv(unset, raising=False) + search: Final = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + with patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}): + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + managed_stores=(), + router=_team_provider_router(), + headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": forwarded_provider_key}, + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {forwarded_provider_key}" + + @pytest.mark.usefixtures("httpx_transport") def test_rag_query_bedrock_search_keeps_the_callers_region_and_credentials_over_team_deployment( client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/rag/ingestion/test_base_ingestion.py b/tests/unit/rag/ingestion/test_base_ingestion.py index 19f10dfc5b9..63a0f0d7da5 100644 --- a/tests/unit/rag/ingestion/test_base_ingestion.py +++ b/tests/unit/rag/ingestion/test_base_ingestion.py @@ -229,7 +229,7 @@ async def test_a_gemini_upstream_error_keeps_its_status_on_the_failed_ingest( @pytest.mark.asyncio -async def test_a_gemini_store_with_no_api_key_anywhere_fails_as_a_bad_request(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_a_gemini_store_with_no_api_key_anywhere_fails_without_a_caller_error_status(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("GEMINI_API_KEY", raising=False) monkeypatch.delenv("GOOGLE_API_KEY", raising=False) ingestion: Final = GeminiRAGIngestion(ingest_options={"vector_store": {"custom_llm_provider": "gemini"}}) @@ -237,7 +237,7 @@ async def test_a_gemini_store_with_no_api_key_anywhere_fails_as_a_bad_request(mo response: Final = await ingestion.ingest(file_data=_TEXT_FILE) assert response["status"] == "failed" - assert response["error_status_code"] == 400 + assert response.get("error_status_code") is None @pytest.mark.asyncio