diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 17f8bd277dd..483e6b54fbb 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -509,6 +509,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_id: str, custom_llm_provider: str, data: dict, + user_api_key_dict: UserAPIKeyAuth | None = None, ) -> tuple[str, dict[str, Any]]: """ Map proxy-facing file IDs to provider file IDs and credentials. @@ -521,6 +522,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): decode_model_from_file_id, get_credentials_for_model, get_original_file_id, + get_team_provider_credentials, ) from litellm.proxy.proxy_server import llm_router @@ -547,21 +549,32 @@ class _PROXY_BatchRateLimiter(CustomLogger): return get_original_file_id(file_id), fetch_kwargs request_model: Final = data.get("model") - if isinstance(request_model, str) and request_model and llm_router is not None: - try: - credentials = get_credentials_for_model( - llm_router=llm_router, - model_id=request_model, - operation_context="batch input file read (rate limiting)", - ) - fetch_kwargs.update(extract_file_access_credentials(credentials)) - fetch_kwargs["model"] = request_model - provider = credentials.get("custom_llm_provider") - if provider: - fetch_kwargs["custom_llm_provider"] = provider - except HTTPException: - pass + if isinstance(request_model, str) and request_model: + if llm_router is not None: + try: + credentials = get_credentials_for_model( + llm_router=llm_router, + model_id=request_model, + operation_context="batch input file read (rate limiting)", + ) + fetch_kwargs.update(extract_file_access_credentials(credentials)) + fetch_kwargs["model"] = request_model + provider = credentials.get("custom_llm_provider") + if provider: + fetch_kwargs["custom_llm_provider"] = provider + except HTTPException: + pass + return file_id, fetch_kwargs + if user_api_key_dict is None: + return file_id, fetch_kwargs + team_credentials: Final = get_team_provider_credentials( + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) + if team_credentials is not None: + fetch_kwargs.update(extract_file_access_credentials(team_credentials)) return file_id, fetch_kwargs async def _reserve_batch_enqueued_tokens( @@ -912,6 +925,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_id=file_id, custom_llm_provider=custom_llm_provider, data=data or {}, + user_api_key_dict=user_api_key_dict, ) # For non-managed files, use the standard litellm.afile_content file_content = await litellm.afile_content( diff --git a/tests/integration/authorization/test_rag_team_provider_credentials.py b/tests/integration/authorization/test_rag_team_provider_credentials.py index d85c20cad35..f863b57e803 100644 --- a/tests/integration/authorization/test_rag_team_provider_credentials.py +++ b/tests/integration/authorization/test_rag_team_provider_credentials.py @@ -46,13 +46,24 @@ GEMINI_API_KEY: Final = "integration-rag-gemini-key" AZURE_API_KEY: Final = "integration-rag-azure-key" JWT_KEY_ID: Final = "integration-rag-team-provider-jwt" WILDCARD_MODEL: Final = "openai/gpt-4o-mini" +BARE_WILDCARD_MODEL: Final = "gpt-4o-mini" OPENAI: Final[Mapping[str, JsonValue]] = MappingProxyType({"custom_llm_provider": "openai"}) DOCUMENT: Final = ("document.txt", b"team provider credentials", "text/plain") +BATCH_INPUT_ROW: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "custom_id": "integration-rag-batch", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": BARE_WILDCARD_MODEL, "messages": [{"role": "user", "content": "batch input"}]}, + } +) +BATCH_INPUT: Final = ("batch.jsonl", (json.dumps(dict(BATCH_INPUT_ROW)) + "\n").encode(), "application/jsonl") 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)(/[^?]*)") STORE_FILES: Final = re.compile(r"^/v1/vector_stores/([^/]+)/files$") +FILE_CONTENT: Final = re.compile(r"^/v1/files/([^/]+)/content$") STORE_SEARCH: Final = re.compile(r"^(?:/v1|/openai)/vector_stores/([^/]+)/search$") GEMINI_UPLOAD: Final = re.compile(r"^/upload/v1beta/fileSearchStores/([^:]+):uploadToFileSearchStore$") @@ -78,13 +89,18 @@ class Call: @dataclass(frozen=True, slots=True) class Tenants: shared_model: str + global_model: str team_a: str team_a_model: str team_a_key: str restricted_key: str viewer_key: str jwt_member: str + team_w: str team_w_key: str + wildcard_bare_key: str + wildcard_prefixed_key: str + wildcard_global_key: str team_s_key: str teamless_key: str credential_name: str @@ -180,6 +196,30 @@ def _file_reply() -> Reply: ) +def _batch_input_reply() -> Reply: + return Reply(status=200, body=BATCH_INPUT[1]) + + +def _batch_reply(request: Request) -> Reply: + body: Final = _json_body(request) + 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, + } + ) + + def _attach_reply(request: Request, store_id: str) -> Reply: return _json_reply( { @@ -292,12 +332,16 @@ def _gemini_reply(request: Request, route: str) -> Reply: def _openai_reply(request: Request, 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) 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 == "/v1/batches": + return _batch_reply(request) if route == "/v1/chat/completions": return _chat_reply(request) if route == "/v1/responses": @@ -398,6 +442,16 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants: api_base=f"{upstream}/shared/v1", team_id=None, ) + global_model: Final = f"openai/integration-rag-global-{uuid.uuid4().hex}" + _deployment( + gateway, + scenario, + name=global_model, + model="openai/gpt-4o-mini", + api_key=SHARED_API_KEY, + api_base=f"{upstream}/shared/v1", + team_id=None, + ) team_a_model: Final = f"integration-rag-team-{uuid.uuid4().hex}" team_a: Final = scenario.team(models=[team_a_model, shared_model]) _deployment( @@ -424,13 +478,18 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants: team_s: Final = scenario.team(models=[shared_model]) return Tenants( shared_model=shared_model, + global_model=global_model, team_a=team_a, team_a_model=team_a_model, team_a_key=scenario.key(team_id=team_a), restricted_key=scenario.key(team_id=team_a, models=[shared_model]), viewer_key=scenario.key(user_id=viewer, team_id=team_a), jwt_member=scenario.member(team_a), + team_w=team_w, team_w_key=scenario.key(team_id=team_w), + 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_key=scenario.key(team_id=team_s), teamless_key=scenario.key(), credential_name=_credential(gateway, scenario, f"{upstream}/credential/v1"), @@ -726,6 +785,90 @@ def test_files_upload_keeps_the_accessible_deployment_fallback( assert rag_rig.signatures() == ((prefix, "POST", "/v1/files", f"Bearer {api_key}"),) +@pytest.mark.parametrize( + ("caller", "model"), + (("wildcard_bare_key", BARE_WILDCARD_MODEL), ("wildcard_prefixed_key", WILDCARD_MODEL)), + ids=("bare_model", "prefixed_model"), +) +def test_key_restricted_to_a_model_under_the_team_wildcard_uses_the_team_deployment( + rag_rig: RagRig, caller: Literal["wildcard_bare_key", "wildcard_prefixed_key"], model: str +) -> None: + key: Final = getattr(rag_rig.tenants, caller) + + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, key)) + assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + rows: Final = _store_rows(store_id) + assert [row["team_id"] for row in rows] == [rag_rig.tenants.team_w], rows + assert WILDCARD_API_KEY not in json.dumps(rows), rows + + response: Final = _query(rag_rig.gateway, key, model, {**OPENAI, "vector_store_id": store_id}) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("wildcard", WILDCARD_API_KEY, store_id), + _chat_signature("wildcard", WILDCARD_API_KEY), + ) + + +def test_key_restricted_to_a_global_model_under_the_team_wildcard_splits_rag_and_chat(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.wildcard_global_key)) + assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + + response: Final = _query( + rag_rig.gateway, tenants.wildcard_global_key, tenants.global_model, {**OPENAI, "vector_store_id": store_id} + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("wildcard", WILDCARD_API_KEY, store_id), + _chat_signature("shared", SHARED_API_KEY), + ) + + +def test_files_upload_by_a_key_restricted_under_the_team_wildcard_uses_the_team_deployment(rag_rig: RagRig) -> None: + response: Final = rag_rig.gateway.request_multipart( + "/v1/files", {"purpose": "assistants"}, {"file": DOCUMENT}, key=rag_rig.tenants.wildcard_bare_key + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == (("wildcard", "POST", "/v1/files", f"Bearer {WILDCARD_API_KEY}"),) + + +@pytest.mark.parametrize( + ("caller", "prefix", "api_key", "reads_input"), + (("team_a_key", "team", TEAM_API_KEY, False), ("wildcard_bare_key", "wildcard", WILDCARD_API_KEY, True)), + ids=("team_key", "restricted_wildcard_key"), +) +def test_batch_create_reads_the_input_and_creates_the_batch_through_the_team_deployment( + rag_rig: RagRig, caller: Literal["team_a_key", "wildcard_bare_key"], prefix: str, api_key: str, reads_input: bool +) -> None: + key: Final = getattr(rag_rig.tenants, caller) + uploaded: Final = rag_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 = rag_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 + assert object_value(created.json())["input_file_id"] == file_id, created.text + bearer: Final = f"Bearer {api_key}" + input_read: Final = ((prefix, "GET", f"/v1/files/{file_id}/content", bearer),) if reads_input else () + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + *input_read, + (prefix, "POST", "/v1/batches", bearer), + ) + + def test_gemini_ingest_never_receives_the_team_openai_key(rag_rig: RagRig) -> None: store_id: Final = _completed_store( rag_rig, _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {"custom_llm_provider": "gemini"}) diff --git a/tests/unit/proxy/hooks/test_batch_rate_limiter.py b/tests/unit/proxy/hooks/test_batch_rate_limiter.py index 2f71b177705..b6de0affc8d 100644 --- a/tests/unit/proxy/hooks/test_batch_rate_limiter.py +++ b/tests/unit/proxy/hooks/test_batch_rate_limiter.py @@ -481,6 +481,53 @@ async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockR assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows) +@pytest.mark.asyncio +async def test_batch_input_without_a_model_is_read_through_the_team_deployment( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm import Router + from litellm.proxy import proxy_server + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "sk-team", + "api_base": "https://team.example/v1", + }, + "model_info": {"id": "team-openai", "team_id": "team-1"}, + } + ] + ), + ) + _, _, batch_limiter = _make_limiters() + rows: Final = _batch_rows(("Hello",)) + environment_route: Final = _serve_file(openai_files, "file-abc123", rows) + team_route: Final = openai_files.get("https://team.example/v1/files/file-abc123/content").mock( + return_value=httpx.Response(200, content="\n".join(json.dumps(row) for row in rows).encode()) + ) + team_key: Final = UserAPIKeyAuth(api_key="sk", team_id="team-1", team_models=["openai/*"], models=[_BATCH_MODEL]) + teamless_key: Final = UserAPIKeyAuth(api_key="sk") + + team_usage: Final = await batch_limiter.count_input_file_usage( + file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=team_key, data={} + ) + teamless_usage: Final = await batch_limiter.count_input_file_usage( + file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=teamless_key, data={} + ) + + assert team_route.call_count == 1 + assert team_route.calls.last.request.headers["authorization"] == "Bearer sk-team" + assert environment_route.call_count == 1 + assert environment_route.calls.last.request.headers["authorization"] == "Bearer sk-test" + assert team_usage.total_tokens == teamless_usage.total_tokens == _token_counter_total(rows) + + @pytest.mark.asyncio async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context( openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch