From 110d4c2ad1e6c6c1a0a8cf52880fc2c930c45357 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 22:05:57 +0000 Subject: [PATCH] test(integration): add passthrough cost cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/scripts/run_integration.sh | 4 + tests/integration/_support/client.py | 7 +- tests/integration/_support/upstream.py | 8 +- tests/integration/contracts.json | 18 + .../cost_calculation/cost_tracking_case.py | 55 ++- .../cost_calculation/cost_tracking_cases.json | 373 ++++++++++++++++++ .../cost_calculation/test_cost_tracking.py | 71 +++- 7 files changed, 512 insertions(+), 24 deletions(-) diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 0d6cdcabd57..08b0281b30f 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -121,6 +121,10 @@ start_proxy() { "LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map" "MODEL_COST_MAP_MIN_MODEL_COUNT=1" "MODEL_COST_MAP_MAX_SHRINK_RATIO=0" + "GEMINI_API_BASE=$INTEGRATION_UPSTREAM_URL" + "ANTHROPIC_API_BASE=$INTEGRATION_UPSTREAM_URL" + "GEMINI_API_KEY=sk-scripted-provider" + "ANTHROPIC_API_KEY=sk-scripted-provider" ) else cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True") diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 5d206c10f3d..0b6771623c0 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -58,13 +58,18 @@ class Gateway: *, key: str | None = None, params: Mapping[str, str] | None = None, + headers: Mapping[str, str] | None = None, ) -> httpx.Response: + request_headers: Final = { + "Authorization": f"Bearer {self.key if key is None else key}", + **(headers or {}), + } return self.client.request( method, path, json=body, params=params, - headers={"Authorization": f"Bearer {self.key if key is None else key}"}, + headers=request_headers, ) def request_multipart( diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 22f6cdcde78..90325b770ab 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -194,9 +194,11 @@ class Provider: async def scripted(self, request: Request) -> Response: segments: Final = tuple(segment for segment in cast(str, request.path_params["path"]).split("/") if segment) - if not segments: - return JSONResponse({"error": "Unknown scenario"}, status_code=404) - scenario_id: Final = segments[0].split(":", 1)[0] + scenario_id: Final = ( + segments[0].split(":", 1)[0] + if segments and self.scenario_store.get(segments[0].split(":", 1)[0]) is not None + else request.headers.get("x-scripted-scenario", "") + ) response: Final = self.scenario_store.get(scenario_id) if response is None: return JSONResponse({"error": "Unknown scenario"}, status_code=404) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index bb6f7399293..e86536e13b5 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -1431,6 +1431,24 @@ ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-messages_cache_read]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-3.1-pro-passthrough-generate_content_priced_via_gemini_key]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-3.1-pro-passthrough-stream_generate_content_priced_via_vertex_key]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-passthrough-messages]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-passthrough-messages_cache_read]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-passthrough-converse]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-passthrough-converse_stream]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" ] }, "browser": { diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 1869a394016..ba64cab81ed 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -144,6 +144,7 @@ class ExactExpected(BaseModel): reasoning_cost: float | None = None tool_usage_cost: float | None = None breakdown_persisted: bool = True + cost_header: bool = True class RecountRates(BaseModel): @@ -180,19 +181,22 @@ class CostTrackingTestCase(BaseModel): name: str covers: str model: str - endpoint: Literal[ - "/v1/chat/completions", - "/v1/responses", - "/v1/messages", - "/v1/embeddings", - "/v1/rerank", - "/v1/completions", - "/v1/moderations", - "/v1/audio/transcriptions", - "/v1/audio/speech", - "/v1/images/generations", - "/v1/images/edits", - ] = "/v1/chat/completions" + endpoint: ( + Literal[ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1/embeddings", + "/v1/rerank", + "/v1/completions", + "/v1/moderations", + "/v1/audio/transcriptions", + "/v1/audio/speech", + "/v1/images/generations", + "/v1/images/edits", + ] + | Annotated[str, Field(pattern=r"^/(gemini|anthropic|bedrock)/")] + ) = "/v1/chat/completions" deployment: Deployment | None = None upload: Upload | None = None request: dict[str, JsonValue] @@ -235,6 +239,17 @@ class CostTrackingTestCase(BaseModel): def base_model(self) -> str | None: return self.deployment.base_model if self.deployment else None + @property + def passthrough_provider(self) -> Literal["gemini", "anthropic", "bedrock"] | None: + provider: Final = self.endpoint.removeprefix("/").split("/", 1)[0] + if provider == "gemini": + return "gemini" + if provider == "anthropic": + return "anthropic" + if provider == "bedrock": + return "bedrock" + return None + class _CasesFile(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -362,6 +377,19 @@ def data_errors() -> tuple[str, ...]: and case.response.status != 200 ) ) + invalid_opt_outs: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, ExactExpected) + and ( + ( + not case.expected.breakdown_persisted + and case.passthrough_provider is None + and case.rates.mode != "image_generation" + ) + or (not case.expected.cost_header and case.passthrough_provider is None) + ) + ) return tuple( message for message in ( @@ -374,6 +402,7 @@ def data_errors() -> tuple[str, ...]: f"failure response statuses are inconsistent: {failure_response_mismatches}" if failure_response_mismatches else None, + f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None, ) if message is not None ) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index be0d62d360a..1179dd47592 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -27460,6 +27460,379 @@ "completion_tokens": 380, "cache_read_cost": 0.00405504 } + }, + { + "name": "gemini-3.1-pro-passthrough-generate_content_priced_via_gemini_key", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gemini/gemini-3.1-pro", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "5fdf6b7dd9b9 summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "scripted answer 5fdf6b7dd9b9" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "usageMetadata": { + "promptTokenCount": 1840, + "candidatesTokenCount": 412, + "totalTokenCount": 2252, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 1840 + } + ] + }, + "modelVersion": "gemini-3.1-pro" + } + }, + "expected": { + "spend": 0.008624, + "input_cost": 0.00368, + "output_cost": 0.004944, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false, + "cost_header": false + }, + "endpoint": "/gemini/v1beta/models/$MODEL:generateContent" + }, + { + "name": "gemini-3.1-pro-passthrough-stream_generate_content_priced_via_vertex_key", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gemini-3.1-pro", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "52b6a80ff038 summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": true, + "stream_options": { + "include_usage": true + }, + "allowed_openai_params": [] + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"scripted answer 52b6a80ff038\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\", \"index\": 0}], \"modelVersion\": \"gemini-3.1-pro\"}", + "data: {\"candidates\": [], \"usageMetadata\": {\"promptTokenCount\": 1840, \"candidatesTokenCount\": 412, \"totalTokenCount\": 2252, \"promptTokensDetails\": [{\"modality\": \"TEXT\", \"tokenCount\": 1840}]}, \"modelVersion\": \"gemini-3.1-pro\"}" + ] + }, + "expected": { + "spend": 0.0090552, + "input_cost": 0.003864, + "output_cost": 0.0051912, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false, + "cost_header": false + }, + "endpoint": "/gemini/v1beta/models/$MODEL:streamGenerateContent?alt=sse" + }, + { + "name": "claude-sonnet-5-passthrough-messages", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/anthropic/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": "summarize the attached material in one line" + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 1840, + "output_tokens": 412 + } + } + }, + "expected": { + "spend": 0.0117, + "input_cost": 0.00552, + "output_cost": 0.00618, + "prompt_tokens": 1840, + "completion_tokens": 412, + "breakdown_persisted": false, + "cost_header": false + } + }, + { + "name": "claude-sonnet-5-passthrough-messages_cache_read", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "claude-sonnet-5", + "endpoint": "/anthropic/v1/messages", + "request": { + "model": "$MODEL", + "max_tokens": 412, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "cached text", + "cache_control": { + "type": "ephemeral" + } + } + ] + } + ] + }, + "response": { + "content_type": "application/json", + "body": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [ + { + "type": "text", + "text": "scripted response" + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 640, + "output_tokens": 380, + "cache_read_input_tokens": 12288 + } + } + }, + "expected": { + "spend": 0.0113064, + "input_cost": 0.0056064, + "output_cost": 0.0057, + "prompt_tokens": 12928, + "completion_tokens": 380, + "cache_read_cost": 0.0036864, + "breakdown_persisted": false, + "cost_header": false + } + }, + { + "name": "anthropic.claude-sonnet-5-v1:0-passthrough-converse", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "9aad4de0556c summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": false, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/json", + "body": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "scripted answer 9aad4de0556c" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 42 + } + } + }, + "expected": { + "spend": 0.01287, + "input_cost": 0.006072, + "output_cost": 0.006798, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + }, + "endpoint": "/bedrock/model/$MODEL/converse" + }, + { + "name": "anthropic.claude-sonnet-5-v1:0-passthrough-converse_stream", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "anthropic.claude-sonnet-5-v1:0", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line." + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a9257967d38a summarize the attached material in one line and name the city weather" + } + ] + } + ], + "stream": true, + "stream_options": { + "include_usage": true + }, + "allowed_openai_params": [] + }, + "response": { + "content_type": "application/vnd.amazon.eventstream", + "events": [ + { + "event_type": "messageStart", + "payload": { + "role": "assistant" + } + }, + { + "event_type": "contentBlockDelta", + "payload": { + "delta": { + "text": "scripted answer a9257967d38a" + }, + "contentBlockIndex": 0 + } + }, + { + "event_type": "contentBlockStop", + "payload": { + "contentBlockIndex": 0 + } + }, + { + "event_type": "messageStop", + "payload": { + "stopReason": "end_turn" + } + }, + { + "event_type": "metadata", + "payload": { + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252 + }, + "metrics": { + "latencyMs": 42 + } + } + } + ] + }, + "expected": { + "spend": 0.01287, + "input_cost": 0.006072, + "output_cost": 0.006798, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + }, + "endpoint": "/bedrock/model/$MODEL/converse-stream" } ] } diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index f6547406827..c2ac6e77a3f 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -12,8 +12,10 @@ import zlib import httpx import pytest +from pydantic import JsonValue from integration._support.client import JSON_OBJECT, Gateway +from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.conftest import ( CostBreakdown, approx_equal, @@ -96,6 +98,16 @@ def _assert_stream_has_no_error(response_text: str) -> None: ), f"stream carried an error event: {parsed}" +def _replace_model(value: JsonValue, model_name: str) -> JsonValue: + if isinstance(value, str): + return value.replace("$MODEL", model_name) + if isinstance(value, list): + return [_replace_model(item, model_name) for item in value] + if isinstance(value, dict): + return {key: _replace_model(item, model_name) for key, item in value.items()} + return value + + def _assert_breakdown( case: CostTrackingTestCase, expected: ExactExpected, @@ -139,12 +151,12 @@ def _assert_breakdown( assert actual_component is not None and approx_equal(actual_component, expected_component), ( f"{case.name}: {field} {actual_component} != expected {expected_component}" ) - if case.response.content_type == "application/json": + if expected.cost_header and case.response.content_type == "application/json": header: Final = response.headers.get(header_name) assert header is not None and approx_equal(float(header), expected_component), ( f"{case.name}: {header_name} {header} != expected {expected_component}" ) - if case.response.content_type == "application/json" and any( + if expected.cost_header and case.response.content_type == "application/json" and any( component is not None for component in ( expected.cache_read_cost, @@ -171,11 +183,51 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) marker: Final = sha256(case.name.encode()).hexdigest()[:12] with gateway.scenario() as scenario: key: Final = scenario.key() - model_name: Final = register_scenario_deployment(scenario, case, marker, key) + passthrough_provider: Final = case.passthrough_provider + scenario_id: Final = f"sc-{marker}-{sha256(key.encode()).hexdigest()[:12]}" + scenario_handle: Final = ( + register_scenario(scenario_id, case.response) + if passthrough_provider in {"gemini", "anthropic"} + else None + ) + if scenario_handle is not None: + scenario.cleanups.callback(delete_scenario, scenario_handle) + model_name: Final = ( + case.model + if passthrough_provider in {"gemini", "anthropic"} + else register_scenario_deployment(scenario, case, marker, key) + ) + request_model: Final = ( + case.model.rsplit("/", 1)[-1] + if passthrough_provider in {"gemini", "anthropic"} + else model_name + ) + request_body: Final = JSON_OBJECT.validate_python( + _replace_model(case.request, request_model) + if passthrough_provider is not None + else {**case.request, "model": model_name} + ) + request_headers: Final = ( + { + "x-pass-x-scripted-scenario": scenario_id, + **( + {"x-goog-api-key": key} + if passthrough_provider == "gemini" + else {} + ), + } + if passthrough_provider is not None + else {} + ) + request_path: Final = ( + case.endpoint.replace("$MODEL", request_model) + if passthrough_provider is not None + else case.endpoint + ) response: Final = ( _multipart_request(gateway, case, model_name, key) if case.upload is not None - else gateway.request("POST", case.endpoint, {**case.request, "model": model_name}, key=key) + else gateway.request("POST", request_path, request_body, key=key, headers=request_headers) ) if isinstance(case.expected, FailureExpected): assert response.status_code == case.expected.failure.status, ( @@ -220,9 +272,14 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) ) elif case.response.content_type == "application/json": header: Final = cast(str | None, response.headers.get("x-litellm-response-cost")) - assert header is not None and approx_equal(float(header), expected.spend), ( - f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" - ) + if expected.cost_header: + assert header is not None and approx_equal(float(header), expected.spend), ( + f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" + ) + elif header is not None: + assert approx_equal(float(header), expected.spend), ( + f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" + ) assert row.spend is not None and approx_equal(row.spend, expected.spend), ( f"{case.name}: spend {row.spend} != expected {expected.spend} " f"(breakdown {row.breakdown.model_dump() if row.breakdown is not None else None})"