diff --git a/litellm/proxy/common_utils/request_tag_mappings.py b/litellm/proxy/common_utils/request_tag_mappings.py index dca269dca16..3673b8916d9 100644 --- a/litellm/proxy/common_utils/request_tag_mappings.py +++ b/litellm/proxy/common_utils/request_tag_mappings.py @@ -1,3 +1,4 @@ +import re from collections.abc import Mapping, Sequence from typing import Final from urllib.parse import unquote_plus @@ -13,13 +14,18 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.types.proxy.request_tag_mappings import RequestTagMapping REQUEST_TAG_MAPPINGS_ADAPTER: Final = TypeAdapter(tuple[RequestTagMapping, ...]) +_PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE: Final = re.compile(r"/[^/]+/v1/(?:batches|files)(?:/.*)?") def is_baggage_mapped_route(route: str) -> bool: from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers return ( - (RouteChecks.is_llm_api_route(route) or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route)) + ( + RouteChecks.is_llm_api_route(route) + or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route) + or _PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE.fullmatch(route) is not None + ) and not RouteChecks.is_unbilled_route(route) and not RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value) ) diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index e08348a074c..4511df58842 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -143,9 +143,27 @@ def _anthropic_reply() -> Reply: ) +def _batch_reply() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"batch_{uuid.uuid4().hex}", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-baggage", + "completion_window": "24h", + "status": "validating", + "created_at": 1, + } + ).encode() + ) + + def _provider(request: Request) -> Reply: if request.target == GUARDRAIL_PATH: return Reply(body=b'{"action":"NONE"}') + if request.target == "/v1/batches": + return _batch_reply() if request.target == "/v1/responses": return _responses_reply() if request.target == "/v1/messages": @@ -543,6 +561,47 @@ def test_required_mapping_rejects_a_request_without_its_baggage_value_before_the assert admitted.status_code == 200, admitted.text +def _batch_body(rig: Rig) -> dict[str, JsonValue]: + return { + "model": rig.chat_model, + "input_file_id": "file-baggage", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + } + + +@pytest.mark.parametrize("provider_prefix", ["litellm_proxy", "hosted_vllm"]) +def test_required_mapping_rejects_provider_prefixed_batches_and_files_before_the_provider( + mapped: Rig, provider_prefix: str +) -> None: + batches: Final = f"/{provider_prefix}/v1/batches" + files: Final = f"/{provider_prefix}/v1/files" + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key(allowed_routes=[batches, files]) + mapped.provider.drain() + rejected_batch: Final = mapped.proxy.request("POST", batches, _batch_body(mapped), key=key) + rejected_file: Final = mapped.proxy.request_multipart( + files, {"purpose": "batch"}, {"file": ("batch.jsonl", b"{}", "application/jsonl")}, key=key + ) + assert rejected_batch.status_code == 400, rejected_batch.text + assert rejected_file.status_code == 400, rejected_file.text + assert "baggage value for tenant" in rejected_batch.text + assert "baggage value for tenant" in rejected_file.text + assert mapped.provider.drain() == () + + +def test_provider_prefixed_batch_spend_row_carries_the_mapped_tags(mapped: Rig) -> None: + route: Final = "/litellm_proxy/v1/batches" + tenant: Final = uuid.uuid4().hex + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key(allowed_routes=[route]) + admitted: Final = mapped.proxy.request( + "POST", route, _batch_body(mapped), key=key, headers={"baggage": f"entrypoint=a,tenant={tenant}"} + ) + assert admitted.status_code == 200, admitted.text + assert _spend_tags(mapped.proxy, key) == {mapped.exact_tag, mapped.tenant_tag(tenant)} + + def test_require_budgeted_tag_rejects_an_unbudgeted_mapped_tag_until_the_tag_has_a_budget(mapped: Rig) -> None: value: Final = uuid.uuid4().hex tag: Final = mapped.budgeted_tag(value) diff --git a/tests/unit/proxy/common_utils/test_request_tag_mappings.py b/tests/unit/proxy/common_utils/test_request_tag_mappings.py index e1644544c46..dfdbb4e5479 100644 --- a/tests/unit/proxy/common_utils/test_request_tag_mappings.py +++ b/tests/unit/proxy/common_utils/test_request_tag_mappings.py @@ -290,6 +290,19 @@ def test_unbilled_and_mcp_routes_skip_baggage_mappings(route): assert not is_baggage_mapped_route(route) +@pytest.mark.parametrize("prefix", ["/litellm_proxy", "/hosted_vllm", "/xai", "/manus", "/deepseek", "/mcp"]) +@pytest.mark.parametrize("path", ["/v1/batches", "/v1/batches/batch_qa/cancel", "/v1/batches/a/b/cancel", "/v1/files"]) +def test_batch_and_file_routes_under_any_provider_prefix_take_baggage_mappings(prefix, path): + assert is_baggage_mapped_route(f"{prefix}{path}") + + +@pytest.mark.parametrize( + "route", ["/hosted_vllm/nested/v1/batches", "/hosted_vllm/v1/batches_extra", "/hosted_vllm/v2/files"] +) +def test_routes_outside_the_provider_prefixed_batch_and_file_handlers_skip_baggage_mappings(route): + assert not is_baggage_mapped_route(route) + + @pytest.mark.parametrize("registered", [True, False]) def test_custom_pass_through_routes_take_baggage_mappings_once_registered(registered): from litellm.proxy.pass_through_endpoints.pass_through_endpoints import _registered_pass_through_routes