From 0fd8a42972de6413c1c4467e80c582fa689885e8 Mon Sep 17 00:00:00 2001 From: Gaurav Pandey <112387553+gaurav-pandey-zocdoc@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:22:03 +0530 Subject: [PATCH 001/349] fix(alerting): clarify budget threshold messages Generated with AI Co-Authored-By: Claude Code --- litellm/integrations/SlackAlerting/slack_alerting.py | 4 ++-- .../SlackAlerting/test_slack_alerting.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 94d734546be..432cbd0917b 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -632,10 +632,10 @@ class SlackAlerting(CustomBatchLogger): event_message += f"Budget Crossed\n Total Budget:`{user_info.max_budget}`" elif percent_left <= SLACK_ALERTING_THRESHOLD_5_PERCENT: event = "threshold_crossed" - event_message += "5% Threshold Crossed " + event_message += "5% or less of budget remaining" elif percent_left <= SLACK_ALERTING_THRESHOLD_15_PERCENT: event = "threshold_crossed" - event_message += "15% Threshold Crossed" + event_message += "15% or less of budget remaining" return event, event_message diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 55e2dcdc270..6eaa3147dc3 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -10,6 +10,7 @@ import pytest import litellm from litellm.caching.caching import DualCache +from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.proxy._types import CallInfo, Litellm_EntityType from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys @@ -46,9 +47,8 @@ class TestSlackAlerting(unittest.TestCase): self.assertEqual(result, -0.2) def test_get_event_and_event_message_max_budget(self): - # Initial setup with no event event = None - event_message = "Test Message: " + event_message = get_budget_alert_type("user_budget").get_event_message() # Test case 1: When spend exceeds max_budget user_info = CallInfo( @@ -63,7 +63,7 @@ class TestSlackAlerting(unittest.TestCase): self.assertEqual(event, "budget_crossed") self.assertTrue("Budget Crossed" in event_message) - # Test case 2: When 5% of max_budget is left + event_message = get_budget_alert_type("user_budget").get_event_message() user_info = CallInfo( max_budget=100.0, spend=95.0, @@ -74,9 +74,9 @@ class TestSlackAlerting(unittest.TestCase): user_info=user_info, event=event, event_message=event_message ) self.assertEqual(event, "threshold_crossed") - self.assertTrue("5% Threshold Crossed" in event_message) + self.assertEqual(event_message, "User Budget: 5% or less of budget remaining") - # Test case 3: When 15% of max_budget is left + event_message = get_budget_alert_type("user_budget").get_event_message() user_info = CallInfo( max_budget=100.0, spend=85.0, @@ -87,7 +87,7 @@ class TestSlackAlerting(unittest.TestCase): user_info=user_info, event=event, event_message=event_message ) self.assertEqual(event, "threshold_crossed") - self.assertTrue("15% Threshold Crossed" in event_message) + self.assertEqual(event_message, "User Budget: 15% or less of budget remaining") def test_get_event_and_event_message_soft_budget(self): # Initial setup with no event From 2286bf3eca414cc24e0a03b008a7a4e6b9647c44 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 18:30:16 +0900 Subject: [PATCH 002/349] fix(router): stamp model_group when retrieving a batch Batch token usage is accounted on the retrieve call, not on create: a provider only reports token counts once the job finishes, so the usage arrives on aretrieve_batch and that is the spend log row the tokens land on. Router.acreate_batch stamps the requested model group into its metadata, but Router.aretrieve_batch never did. A batch is retrieved by id, so the request carries no model, and the router fans the lookup out over its deployments - leaving model_group unset on the one record that carries the tokens. /global/activity/model groups the spend logs by model_group, so every batch's tokens were bucketed under an empty group. Stamp the model group inside the per-deployment retrieve attempt, preferring an explicitly requested group and otherwise using the model_name of the deployment that answered, which is unambiguous even when the request named no model. An existing model_group in the metadata is left untouched, so nothing that already resolves a group changes. Scope is limited to aretrieve_batch: acompletion, aresponses and acreate_batch logging are untouched, and cost/spend attribution by model is unchanged. Signed-off-by: mynkyu --- litellm/router.py | 9 ++ .../test_router_batch_retrieve_model_group.py | 118 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py diff --git a/litellm/router.py b/litellm/router.py index 3f450661946..c6c25b6be17 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6137,6 +6137,8 @@ class Router: """ try: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) + requested_model_group: Final = model + metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch") if model is not None: filtered_model_list: ( list[DeploymentTypedDict] | list[dict] | dict | None @@ -6173,6 +6175,13 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) + ## STAMP THE MODEL GROUP FOR SPEND TRACKING ## + # A batch is retrieved by id, so the request carries no model group of its + # own - only the deployment that answered knows it. Batch token usage lands + # on this retrieve call (the provider reports counts once the job finishes), + # so without this the tokens are logged under an empty model_group. + model_group: Final = requested_model_group or model_name["model_name"] + new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) data.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py new file mode 100644 index 00000000000..ef8a23e4917 --- /dev/null +++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py @@ -0,0 +1,118 @@ +""" +model_group attribution on router batch retrieval. + +Batch token usage is accounted on the *retrieve* call, not on create: the +provider only knows the token counts once the job finishes, so +`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the +spend log tokens land on. + +`aretrieve_batch` is addressed by batch_id, so the request carries no model, +and the router fans the lookup out across its deployments. These tests lock +that the winning deployment's model group is stamped on the emitted +StandardLoggingPayload, so `/global/activity/model` - which groups the spend +logs by `model_group` - can attribute those tokens instead of bucketing every +batch under "". +""" + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +import litellm.batches.main as bm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import LiteLLMBatch, Usage + +MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev" +DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite" + + +class _PayloadCollector(CustomLogger): + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs.get("standard_logging_object")) + + +@pytest.fixture +def router(): + return Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": { + "model": DEPLOYMENT_MODEL, + "vertex_project": "fake-project", + "vertex_location": "us-central1", + "vertex_credentials": "fake-creds", + }, + } + ] + ) + + +@pytest.fixture +def collector(): + logger = _PayloadCollector() + previous = litellm.callbacks + litellm.callbacks = [logger] + try: + yield logger + finally: + litellm.callbacks = previous + + +@pytest.fixture +def vertex_retrieve(): + """Mock the vertex provider seam - the only real network boundary.""" + batch = LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-1", + object="batch", + status="completed", + usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), + ) + seam = MagicMock(name="vertex_ai_batches_instance") + seam.retrieve_batch.return_value = batch + with patch.object(bm, "vertex_ai_batches_instance", seam): + yield seam + + +async def _collected_payload(collector) -> dict: + for _ in range(50): # the success handler runs as a background task + payloads = [p for p in collector.payloads if p is not None] + if payloads: + return payloads[-1] + await asyncio.sleep(0.05) + raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}") + + +@pytest.mark.asyncio +async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve): + """ + The proxy retrieves a managed batch by id only - no `model` in the request. + The router fans out over its deployments, so the model group is only known + from the deployment that answered. + """ + response = await router.aretrieve_batch(batch_id="batch-1") + + assert response.usage.total_tokens == 1200 + payload = await _collected_payload(collector) + assert payload["model"] == DEPLOYMENT_MODEL + assert payload["model_group"] == MODEL_GROUP + + +@pytest.mark.asyncio +async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve): + """An explicitly requested model group is what gets logged.""" + await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1") + + payload = await _collected_payload(collector) + assert payload["model_group"] == MODEL_GROUP From e630f21d16b10b78e22c28a974dee73009749167 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 19:01:18 +0900 Subject: [PATCH 003/349] test: fake the provider at the HTTP boundary in the batch model_group test The test-quality gate flagged the first version for patching an SDK internal (litellm.batches.main.vertex_ai_batches_instance) and for writing litellm.callbacks directly. Drive an openai-compatible deployment through respx instead, so the retrieve call and the usage accounting that reads the completed batch's output file both run for real, and install the collector with monkeypatch so nothing leaks into the next test. Signed-off-by: mynkyu --- .../test_router_batch_retrieve_model_group.py | 146 +++++++++++------- 1 file changed, 89 insertions(+), 57 deletions(-) diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py index ef8a23e4917..b99ec50e041 100644 --- a/tests/test_litellm/test_router_batch_retrieve_model_group.py +++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py @@ -1,35 +1,81 @@ """ model_group attribution on router batch retrieval. -Batch token usage is accounted on the *retrieve* call, not on create: the -provider only knows the token counts once the job finishes, so -`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the -spend log tokens land on. +Batch token usage is accounted on the *retrieve* call, not on create: a provider +only reports token counts once the job finishes, so the usage is read off the +completed batch's output file during retrieve logging and that is the spend log +row the tokens land on. -`aretrieve_batch` is addressed by batch_id, so the request carries no model, -and the router fans the lookup out across its deployments. These tests lock -that the winning deployment's model group is stamped on the emitted -StandardLoggingPayload, so `/global/activity/model` - which groups the spend -logs by `model_group` - can attribute those tokens instead of bucketing every -batch under "". +A batch is retrieved by id, so the request carries no model and the router fans +the lookup out across its deployments. These tests lock that the answering +deployment's model group is stamped on the emitted StandardLoggingPayload, so +`/global/activity/model` - which groups the spend logs by `model_group` - can +attribute those tokens instead of bucketing every batch under "". + +The provider is faked at the HTTP boundary, so the whole retrieve + usage +accounting path runs for real. """ import asyncio -from unittest.mock import MagicMock, patch +import json +import httpx import pytest +import respx import litellm -import litellm.batches.main as bm from litellm import Router from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import LiteLLMBatch, Usage -MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev" -DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite" +MODEL_GROUP = "gemini-batch-group" +DEPLOYMENT_MODEL = "openai/gpt-4o-mini" +API_BASE = "http://localhost:4001/v1" +BATCH_ID = "batch-1" +ROWS = 2 +TOKENS_PER_ROW = 600 + +COMPLETED_BATCH = { + "id": BATCH_ID, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-1", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-1", + "error_file_id": None, + "created_at": 0, + "completed_at": 1, + "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0}, + "metadata": None, +} + +OUTPUT_JSONL = "\n".join( + json.dumps( + { + "id": f"req-{row}", + "custom_id": f"row-{row}", + "response": { + "status_code": 200, + "body": { + "id": f"chatcmpl-{row}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW}, + }, + }, + } + ) + for row in range(ROWS) +) class _PayloadCollector(CustomLogger): + """Captures the StandardLoggingPayload the spend log is built from.""" + def __init__(self): super().__init__() self.payloads = [] @@ -37,6 +83,14 @@ class _PayloadCollector(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.payloads.append(kwargs.get("standard_logging_object")) + async def retrieve_batch_payload(self) -> dict: + for _ in range(100): # the success handler runs as a background task + for payload in self.payloads: + if payload and payload.get("call_type") == "aretrieve_batch": + return payload + await asyncio.sleep(0.05) + raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") + @pytest.fixture def router(): @@ -46,9 +100,8 @@ def router(): "model_name": MODEL_GROUP, "litellm_params": { "model": DEPLOYMENT_MODEL, - "vertex_project": "fake-project", - "vertex_location": "us-central1", - "vertex_credentials": "fake-creds", + "api_base": API_BASE, + "api_key": "sk-fake", }, } ] @@ -56,63 +109,42 @@ def router(): @pytest.fixture -def collector(): +def collector(monkeypatch): logger = _PayloadCollector() - previous = litellm.callbacks - litellm.callbacks = [logger] - try: - yield logger - finally: - litellm.callbacks = previous + monkeypatch.setattr(litellm, "callbacks", [logger]) + return logger @pytest.fixture -def vertex_retrieve(): - """Mock the vertex provider seam - the only real network boundary.""" - batch = LiteLLMBatch( - id="batch-1", - completion_window="24h", - created_at=0, - endpoint="/v1/chat/completions", - input_file_id="file-1", - object="batch", - status="completed", - usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), - ) - seam = MagicMock(name="vertex_ai_batches_instance") - seam.retrieve_batch.return_value = batch - with patch.object(bm, "vertex_ai_batches_instance", seam): - yield seam - - -async def _collected_payload(collector) -> dict: - for _ in range(50): # the success handler runs as a background task - payloads = [p for p in collector.payloads if p is not None] - if payloads: - return payloads[-1] - await asyncio.sleep(0.05) - raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}") +def provider(): + """Fake the provider at the HTTP boundary: the completed batch plus the + output file the usage accounting reads.""" + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH)) + respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL)) + yield respx_mock @pytest.mark.asyncio -async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve): +async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider): """ The proxy retrieves a managed batch by id only - no `model` in the request. The router fans out over its deployments, so the model group is only known from the deployment that answered. """ - response = await router.aretrieve_batch(batch_id="batch-1") + response = await router.aretrieve_batch(batch_id=BATCH_ID) - assert response.usage.total_tokens == 1200 - payload = await _collected_payload(collector) + assert response.id == BATCH_ID + payload = await collector.retrieve_batch_payload() + assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW assert payload["model"] == DEPLOYMENT_MODEL assert payload["model_group"] == MODEL_GROUP @pytest.mark.asyncio -async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve): +async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider): """An explicitly requested model group is what gets logged.""" - await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1") + await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID) - payload = await _collected_payload(collector) + payload = await collector.retrieve_batch_payload() assert payload["model_group"] == MODEL_GROUP From df6990c7127a0c30c77b96323e8311d9200a23be Mon Sep 17 00:00:00 2001 From: mynkyu Date: Sun, 6 Sep 2026 10:10:07 +0900 Subject: [PATCH 004/349] test: move the batch model_group regression into test_router.py CLAUDE.md asks bug fixes to extend the existing mapped test file rather than add a new one, and tests/test_litellm/test_router.py already covers Router.aretrieve_batch. Fold the two cases in next to that coverage and drop the standalone file. The helpers are prefixed so they read unambiguously in a shared file, and the respx context stays open while the payload is awaited, since the usage accounting reads the batch's output file from the success handler. Signed-off-by: mynkyu --- tests/test_litellm/test_router.py | 153 ++++++++++++++++++ .../test_router_batch_retrieve_model_group.py | 150 ----------------- 2 files changed, 153 insertions(+), 150 deletions(-) delete mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7c044310e14..9b146092927 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -958,6 +958,159 @@ async def test_arouter_aretrieve_batch(): assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base" +# --------------------------------------------------------------------------- +# Batch retrieval has to attribute its tokens to a model group. +# +# Batch token usage is accounted on the *retrieve* call, not on create: a +# provider only reports token counts once the job finishes, so the usage is read +# off the completed batch's output file during retrieve logging, and that is the +# spend log row the tokens land on. A batch is retrieved by id, so the request +# carries no model and the router fans the lookup out across its deployments - +# the group of the deployment that answered is the only one there is to stamp. +# Leaving it unset files every batch's tokens under an empty model_group, which +# is what /global/activity/model groups the spend logs by. +# +# The provider is faked at the HTTP boundary, so the retrieve call and the usage +# accounting that reads the output file both run for real. +# --------------------------------------------------------------------------- + +_BATCH_GROUP = "gemini-batch-group" +_BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini" +_BATCH_API_BASE = "http://localhost:4001/v1" +_BATCH_ID = "batch-1" +_BATCH_ROWS = 2 +_BATCH_TOKENS_PER_ROW = 600 + +_BATCH_COMPLETED = { + "id": _BATCH_ID, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-1", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-1", + "error_file_id": None, + "created_at": 0, + "completed_at": 1, + "request_counts": {"total": _BATCH_ROWS, "completed": _BATCH_ROWS, "failed": 0}, + "metadata": None, +} + +_BATCH_OUTPUT_JSONL = "\n".join( + json.dumps( + { + "id": f"req-{row}", + "custom_id": f"row-{row}", + "response": { + "status_code": 200, + "body": { + "id": f"chatcmpl-{row}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": 500, + "completion_tokens": 100, + "total_tokens": _BATCH_TOKENS_PER_ROW, + }, + }, + }, + } + ) + for row in range(_BATCH_ROWS) +) + + +class _BatchPayloadCollector(CustomLogger): + """Captures the StandardLoggingPayload the spend log row is built from.""" + + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs.get("standard_logging_object")) + + async def retrieve_batch_payload(self): + for _ in range(100): # the success handler runs as a background task + for payload in self.payloads: + if payload and payload.get("call_type") == "aretrieve_batch": + return payload + await asyncio.sleep(0.05) + raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") + + +def _batch_model_group_router(): + return litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + } + ] + ) + + +def _mock_batch_provider(respx_mock): + """The completed batch, plus the output file the usage accounting reads.""" + respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(200, json=_BATCH_COMPLETED) + ) + respx_mock.get(f"{_BATCH_API_BASE}/files/file-out-1/content").mock( + return_value=httpx.Response(200, text=_BATCH_OUTPUT_JSONL) + ) + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypatch: pytest.MonkeyPatch): + """ + The proxy retrieves a managed batch by id only - no `model` in the request. + The router fans out over its deployments, so the model group is only known + from the deployment that answered. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = _batch_model_group_router() + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + # the usage accounting reads the output file from the success handler, + # so the provider has to stay faked until that payload lands + payload = await collector.retrieve_batch_payload() + + assert response.id == _BATCH_ID + assert payload["total_tokens"] == _BATCH_ROWS * _BATCH_TOKENS_PER_ROW + assert payload["model"] == _BATCH_DEPLOYMENT_MODEL + assert payload["model_group"] == _BATCH_GROUP + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(monkeypatch: pytest.MonkeyPatch): + """An explicitly requested model group is what gets logged.""" + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = _batch_model_group_router() + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + await router.aretrieve_batch(model=_BATCH_GROUP, batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + + assert payload["model_group"] == _BATCH_GROUP + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py deleted file mode 100644 index b99ec50e041..00000000000 --- a/tests/test_litellm/test_router_batch_retrieve_model_group.py +++ /dev/null @@ -1,150 +0,0 @@ -""" -model_group attribution on router batch retrieval. - -Batch token usage is accounted on the *retrieve* call, not on create: a provider -only reports token counts once the job finishes, so the usage is read off the -completed batch's output file during retrieve logging and that is the spend log -row the tokens land on. - -A batch is retrieved by id, so the request carries no model and the router fans -the lookup out across its deployments. These tests lock that the answering -deployment's model group is stamped on the emitted StandardLoggingPayload, so -`/global/activity/model` - which groups the spend logs by `model_group` - can -attribute those tokens instead of bucketing every batch under "". - -The provider is faked at the HTTP boundary, so the whole retrieve + usage -accounting path runs for real. -""" - -import asyncio -import json - -import httpx -import pytest -import respx - -import litellm -from litellm import Router -from litellm.integrations.custom_logger import CustomLogger - -MODEL_GROUP = "gemini-batch-group" -DEPLOYMENT_MODEL = "openai/gpt-4o-mini" -API_BASE = "http://localhost:4001/v1" -BATCH_ID = "batch-1" -ROWS = 2 -TOKENS_PER_ROW = 600 - -COMPLETED_BATCH = { - "id": BATCH_ID, - "object": "batch", - "endpoint": "/v1/chat/completions", - "errors": None, - "input_file_id": "file-in-1", - "completion_window": "24h", - "status": "completed", - "output_file_id": "file-out-1", - "error_file_id": None, - "created_at": 0, - "completed_at": 1, - "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0}, - "metadata": None, -} - -OUTPUT_JSONL = "\n".join( - json.dumps( - { - "id": f"req-{row}", - "custom_id": f"row-{row}", - "response": { - "status_code": 200, - "body": { - "id": f"chatcmpl-{row}", - "object": "chat.completion", - "model": "gpt-4o-mini", - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} - ], - "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW}, - }, - }, - } - ) - for row in range(ROWS) -) - - -class _PayloadCollector(CustomLogger): - """Captures the StandardLoggingPayload the spend log is built from.""" - - def __init__(self): - super().__init__() - self.payloads = [] - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.payloads.append(kwargs.get("standard_logging_object")) - - async def retrieve_batch_payload(self) -> dict: - for _ in range(100): # the success handler runs as a background task - for payload in self.payloads: - if payload and payload.get("call_type") == "aretrieve_batch": - return payload - await asyncio.sleep(0.05) - raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") - - -@pytest.fixture -def router(): - return Router( - model_list=[ - { - "model_name": MODEL_GROUP, - "litellm_params": { - "model": DEPLOYMENT_MODEL, - "api_base": API_BASE, - "api_key": "sk-fake", - }, - } - ] - ) - - -@pytest.fixture -def collector(monkeypatch): - logger = _PayloadCollector() - monkeypatch.setattr(litellm, "callbacks", [logger]) - return logger - - -@pytest.fixture -def provider(): - """Fake the provider at the HTTP boundary: the completed batch plus the - output file the usage accounting reads.""" - with respx.mock(assert_all_called=True) as respx_mock: - respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH)) - respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL)) - yield respx_mock - - -@pytest.mark.asyncio -async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider): - """ - The proxy retrieves a managed batch by id only - no `model` in the request. - The router fans out over its deployments, so the model group is only known - from the deployment that answered. - """ - response = await router.aretrieve_batch(batch_id=BATCH_ID) - - assert response.id == BATCH_ID - payload = await collector.retrieve_batch_payload() - assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW - assert payload["model"] == DEPLOYMENT_MODEL - assert payload["model_group"] == MODEL_GROUP - - -@pytest.mark.asyncio -async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider): - """An explicitly requested model group is what gets logged.""" - await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID) - - payload = await collector.retrieve_batch_payload() - assert payload["model_group"] == MODEL_GROUP From 128cb114bdac6b8cf41a9d689f0a573a2e27eced Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:45:47 -0700 Subject: [PATCH 005/349] style: trim comments on batch retrieve model group stamp --- litellm/router.py | 7 ++----- tests/test_litellm/test_router.py | 16 ---------------- 2 files changed, 2 insertions(+), 21 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index c6c25b6be17..20c9018abee 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6175,11 +6175,8 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) - ## STAMP THE MODEL GROUP FOR SPEND TRACKING ## - # A batch is retrieved by id, so the request carries no model group of its - # own - only the deployment that answered knows it. Batch token usage lands - # on this retrieve call (the provider reports counts once the job finishes), - # so without this the tokens are logged under an empty model_group. + # A batch is retrieved by id, so only the deployment that answered knows the + # group, and batch token usage is logged on this retrieve call. model_group: Final = requested_model_group or model_name["model_name"] new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9b146092927..258d1c973d0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -958,22 +958,6 @@ async def test_arouter_aretrieve_batch(): assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base" -# --------------------------------------------------------------------------- -# Batch retrieval has to attribute its tokens to a model group. -# -# Batch token usage is accounted on the *retrieve* call, not on create: a -# provider only reports token counts once the job finishes, so the usage is read -# off the completed batch's output file during retrieve logging, and that is the -# spend log row the tokens land on. A batch is retrieved by id, so the request -# carries no model and the router fans the lookup out across its deployments - -# the group of the deployment that answered is the only one there is to stamp. -# Leaving it unset files every batch's tokens under an empty model_group, which -# is what /global/activity/model groups the spend logs by. -# -# The provider is faked at the HTTP boundary, so the retrieve call and the usage -# accounting that reads the output file both run for real. -# --------------------------------------------------------------------------- - _BATCH_GROUP = "gemini-batch-group" _BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini" _BATCH_API_BASE = "http://localhost:4001/v1" From 33d89c9814641a34cb66d357e2cc3a403677a06d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:55:11 -0700 Subject: [PATCH 006/349] style: drop redundant comments per repo comment policy --- litellm/router.py | 3 +-- tests/test_litellm/test_router.py | 3 --- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 20c9018abee..e85ca1bd7a8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6175,8 +6175,7 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) - # A batch is retrieved by id, so only the deployment that answered knows the - # group, and batch token usage is logged on this retrieve call. + # Batch token usage is logged on this retrieve call, not on create. model_group: Final = requested_model_group or model_name["model_name"] new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 258d1c973d0..dc51339cf55 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1009,8 +1009,6 @@ _BATCH_OUTPUT_JSONL = "\n".join( class _BatchPayloadCollector(CustomLogger): - """Captures the StandardLoggingPayload the spend log row is built from.""" - def __init__(self): super().__init__() self.payloads = [] @@ -1043,7 +1041,6 @@ def _batch_model_group_router(): def _mock_batch_provider(respx_mock): - """The completed batch, plus the output file the usage accounting reads.""" respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( return_value=httpx.Response(200, json=_BATCH_COMPLETED) ) From 01bdfb34aa5ed88320dbd1c9f231876df820ca29 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:24:20 -0700 Subject: [PATCH 007/349] chore: drop redundant comments in aretrieve_batch router tests --- tests/test_litellm/test_router.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index dc51339cf55..bfa9ea7e3d1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1017,7 +1017,7 @@ class _BatchPayloadCollector(CustomLogger): self.payloads.append(kwargs.get("standard_logging_object")) async def retrieve_batch_payload(self): - for _ in range(100): # the success handler runs as a background task + for _ in range(100): for payload in self.payloads: if payload and payload.get("call_type") == "aretrieve_batch": return payload @@ -1065,8 +1065,6 @@ async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypa with respx.mock(assert_all_called=True) as respx_mock: _mock_batch_provider(respx_mock) response = await router.aretrieve_batch(batch_id=_BATCH_ID) - # the usage accounting reads the output file from the success handler, - # so the provider has to stay faked until that payload lands payload = await collector.retrieve_batch_payload() assert response.id == _BATCH_ID From 9acf09f60d4f917053e1bfb9d493dce3cdd2771a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:07:55 -0700 Subject: [PATCH 008/349] fix(router): keep batch retrieves out of the per-minute tpm/rpm counters Stamping model_group let both router deployment callbacks past their `model_group is None` early return for batch retrieves. A batch reports the whole job's token total on retrieve and reports it again on every poll of the finished batch, so those tokens are not load in the current minute: three polls of one completed 1,200 token batch pushed a tpm:1000 deployment to 3,600. The fan-out also probed unrelated deployments, adding an rpm tick to each. --- litellm/router.py | 5 ++ litellm/router_utils/batch_utils.py | 18 +++++++ tests/test_litellm/test_router.py | 76 +++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index e85ca1bd7a8..9dd267d7560 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -124,6 +124,7 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, + is_batch_retrieve_call_type, replace_model_in_jsonl, should_replace_model_in_jsonl, ) @@ -7878,6 +7879,8 @@ class Router: # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object is None") @@ -8117,6 +8120,8 @@ class Router: """ Update RPM usage for a deployment """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return deployment_name: Final = kwargs["litellm_params"]["metadata"].get( "deployment", None ) # handles wildcard routes - by giving the original name sent to `litellm.completion` diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index ccb6ad95519..6e110b586fb 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -5,6 +5,7 @@ from typing import Final from litellm._logging import verbose_logger from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose +from litellm.types.utils import CallTypes class InMemoryFile(io.BytesIO): @@ -170,3 +171,20 @@ def _get_router_metadata_variable_name(function_name: str | None) -> str: return "litellm_metadata" else: return "metadata" + + +BATCH_RETRIEVE_CALL_TYPES: Final = frozenset( + { + CallTypes.aretrieve_batch.value, + CallTypes.retrieve_batch.value, + } +) + + +def is_batch_retrieve_call_type(call_type: object) -> bool: + """ + A batch retrieve reports the whole job's token usage, which the provider spent + asynchronously over the life of the batch, and reports it again on every poll of the + finished batch. Per-minute usage counters must not be fed from it. + """ + return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bfa9ea7e3d1..d8045722998 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1090,6 +1090,82 @@ async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(m assert payload["model_group"] == _BATCH_GROUP +_UNRELATED_BATCH_GROUP = "unrelated-batch-group" +_UNRELATED_BATCH_API_BASE = "http://localhost:4002/v1" + +_BATCH_NOT_FOUND = { + "error": { + "message": f"No batch found with id '{_BATCH_ID}'.", + "type": "invalid_request_error", + "code": "batch_not_found", + } +} + + +async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:")) + if keys: + return keys + await asyncio.sleep(0.05) + return [] + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch): + """ + A batch reports the whole job's tokens on retrieve, and reports them again on every + poll of the finished batch, so they are not a measure of load in the current minute. + The fan-out also probes deployments the caller never named. Neither may reach the + per-minute tpm/rpm counters that gate live traffic. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "batch-dep"}, + "tpm": 1000, + "rpm": 10, + }, + { + "model_name": _UNRELATED_BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _UNRELATED_BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "unrelated-dep"}, + "tpm": 1000, + "rpm": 10, + }, + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + usage_keys = await _router_usage_keys(router) + + assert response.id == _BATCH_ID + assert payload["model_group"] == _BATCH_GROUP + assert usage_keys == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ From 58c3d04733f2bebfbc15e8f1f6dd702a37c6e2f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:24:47 -0700 Subject: [PATCH 009/349] test: cover is_batch_retrieve_call_type in router batch utils --- .../router_unit_tests/test_router_batch_utils.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index c9f19731372..e274ac61a01 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -317,3 +317,18 @@ def test_replace_model_in_jsonl_with_embedded_newlines(): == "This is a message\nwith multiple\nlines" ) assert result_json["custom_id"] == "test123" + + +def test_is_batch_retrieve_call_type_matches_only_batch_retrieves(): + from litellm.router_utils.batch_utils import is_batch_retrieve_call_type + from litellm.types.utils import CallTypes + + assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True + assert is_batch_retrieve_call_type(CallTypes.retrieve_batch.value) is True + + for call_type in CallTypes: + if call_type in (CallTypes.aretrieve_batch, CallTypes.retrieve_batch): + continue + assert is_batch_retrieve_call_type(call_type.value) is False + + assert is_batch_retrieve_call_type(None) is False From ad2afe5e6568b69389c2258680274098b9191b6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:41:07 -0700 Subject: [PATCH 010/349] style(router): drop the inline comment on the batch retrieve stamp --- litellm/router.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 9dd267d7560..2397c6b2fc7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6176,7 +6176,6 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) - # Batch token usage is logged on this retrieve call, not on create. model_group: Final = requested_model_group or model_name["model_name"] new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) From 63ea19743373fa1dd87081a66d64e5f580212f77 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:02:53 -0700 Subject: [PATCH 011/349] fix(router): keep batch retrieves out of routing strategy state Stamping model_group on a batch retrieve routed the whole batch job's token usage into the per-model-group counters that usage-based, latency-based, cost-based and least-busy routing read, so polling a finished batch could exhaust a group's TPM or RPM window and lock live chat traffic out with RouterRateLimitError. Polling also drove the least-busy in-flight counts negative once per poll per deployment, which pinned chat to whichever deployment had been polled most. The strategy callbacks now skip batch retrieve call types, so a retrieve still lands in spend logs under its model group while the numbers that pick a deployment for the next chat request stay driven by live traffic only. --- litellm/router.py | 3 +- litellm/router_strategy/least_busy.py | 11 ++++ litellm/router_strategy/lowest_cost.py | 5 ++ litellm/router_strategy/lowest_latency.py | 7 ++ litellm/router_strategy/lowest_tpm_rpm.py | 5 ++ litellm/router_utils/batch_utils.py | 3 +- tests/test_litellm/test_router.py | 79 +++++++++++++++++++++++ 7 files changed, 111 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2397c6b2fc7..8f913d96463 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6177,7 +6177,8 @@ class Router: function_name="aretrieve_batch", ) model_group: Final = requested_model_group or model_name["model_name"] - new_kwargs[metadata_variable_name].setdefault("model_group", model_group) + if not new_kwargs[metadata_variable_name].get("model_group"): + new_kwargs[metadata_variable_name]["model_group"] = model_group new_kwargs.pop("custom_llm_provider", None) data.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..e93288fd9fe 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -11,6 +11,7 @@ from typing import Final from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type class LeastBusyLoggingHandler(CustomLogger): @@ -27,6 +28,8 @@ class LeastBusyLoggingHandler(CustomLogger): Caching based on model group. """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -48,6 +51,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -76,6 +81,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass def log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -103,6 +110,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -131,6 +140,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index b927df0c438..aaad6186484 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -8,6 +8,7 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type class LowestCostLoggingHandler(CustomLogger): @@ -19,6 +20,8 @@ class LowestCostLoggingHandler(CustomLogger): self.router_cache = router_cache def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update usage on success @@ -96,6 +99,8 @@ class LowestCostLoggingHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update cost usage on success diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index a1b67eaeaf9..598ca1227ec 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -9,6 +9,7 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -35,6 +36,8 @@ class LowestLatencyLoggingHandler(CustomLogger): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update latency usage on success @@ -167,6 +170,8 @@ class LowestLatencyLoggingHandler(CustomLogger): """ Check if Timeout Error, if timeout set deployment latency -> 100 """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: metadata_field: Final = self._select_metadata_field(kwargs) _exception: Final = kwargs.get("exception", None) @@ -221,6 +226,8 @@ class LowestLatencyLoggingHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update latency usage on success diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 31c4b1d7e3f..d4abf1f8f70 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -8,6 +8,7 @@ from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose @@ -27,6 +28,8 @@ class LowestTPMLoggingHandler(CustomLogger): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success @@ -79,6 +82,8 @@ class LowestTPMLoggingHandler(CustomLogger): verbose_router_logger.debug(traceback.format_exc()) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 6e110b586fb..be20c358202 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -185,6 +185,7 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: """ A batch retrieve reports the whole job's token usage, which the provider spent asynchronously over the life of the batch, and reports it again on every poll of the - finished batch. Per-minute usage counters must not be fed from it. + finished batch. The counters that measure live traffic, per-minute rate limits and the + routing strategies' own state, must not be fed from it. """ return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d8045722998..836085c1f5d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1166,6 +1166,85 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m assert usage_keys == [] +_ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") + + +async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + keys = sorted( + key + for key in router.cache.in_memory_cache.cache_dict + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) + ) + if keys: + return keys + await asyncio.sleep(0.05) + return [] + + +def _batch_fan_out_router(routing_strategy: str): + return litellm.Router( + routing_strategy=routing_strategy, + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "batch-dep"}, + }, + { + "model_name": _UNRELATED_BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _UNRELATED_BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "unrelated-dep"}, + }, + ], + ) + + +@pytest.mark.parametrize( + "routing_strategy", + ["usage-based-routing", "latency-based-routing", "cost-based-routing", "least-busy"], +) +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( + monkeypatch: pytest.MonkeyPatch, routing_strategy: str +): + """ + Every routing strategy picks a deployment from what recent live traffic did. + A batch retrieve reports the whole job on every poll and probes deployments the + caller never named, so polling a finished batch must not move the numbers that + decide where the next chat request goes. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + monkeypatch.setattr(litellm, "input_callback", []) + router = _batch_fan_out_router(routing_strategy) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + for _ in range(3): + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + await collector.retrieve_batch_payload() + strategy_keys = await _router_strategy_keys(router) + + assert response.id == _BATCH_ID + assert strategy_keys == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ From 828a02f78f2c974f6458237c6fde1128b01b26bb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 03:03:45 -0700 Subject: [PATCH 012/349] fix(batches): stamp the model group on the proxy's model-encoded retrieve path The model-encoded batch id path calls the SDK directly, so the router never labels it. Stamp the decoded group into the request's litellm_metadata, and guard usage-based-routing-v2 the same way the other strategies already are. --- litellm/proxy/batches_endpoints/endpoints.py | 21 +++++++--- litellm/router_strategy/lowest_tpm_rpm_v2.py | 5 +++ .../proxy/batches_endpoints/test_endpoints.py | 40 ++++++++++++++++++- tests/test_litellm/test_router.py | 26 +++++++----- 4 files changed, 76 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..b6a8b421e02 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -52,6 +52,20 @@ from litellm.types.llms.openai import LiteLLMBatchCreateRequest router: Final = APIRouter() +def _litellm_metadata_of(data: dict) -> dict: + """The request's litellm_metadata mapping, created on the request when it carries none. + + The success handler reads this mapping, so a flag or a model group set here has to live + inside it rather than beside it. + """ + existing: Final = data.get("litellm_metadata") + if isinstance(existing, dict): + return existing + created: Final = {} # mutable-ok: the logging layer copies and extends this mapping, so it cannot be a read-only view + data["litellm_metadata"] = created + return created + + def _raise_not_found_when_openai_fallback_unservable( requested_provider: "str | None", data: Mapping[str, object], @@ -531,11 +545,7 @@ async def retrieve_batch( poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active() if poller_owns_accounting: - litellm_metadata = data.get("litellm_metadata") - if not isinstance(litellm_metadata, dict): - litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend - data["litellm_metadata"] = litellm_metadata - litellm_metadata["batch_ignore_default_logging"] = True + _litellm_metadata_of(data)["batch_ignore_default_logging"] = True # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info @@ -558,6 +568,7 @@ async def retrieve_batch( # so litellm.aretrieve_batch can load BedrockBatchesConfig. Without # it the call falls into the legacy provider switch and 400s. data["model"] = model_from_id + _litellm_metadata_of(data).setdefault("model_group", model_from_id) # Retrieve batch using model credentials response = await litellm.aretrieve_batch( diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 665ff69ab47..a2acce5fcb5 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -210,6 +211,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): return deployment # don't fail calls if eg. redis fails to connect def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success @@ -250,6 +253,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM usage on success diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..c6df8f2ffcf 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1233,9 +1233,11 @@ async def call_retrieve( user: Optional[UserAPIKeyAuth] = None, headers: Optional[Dict[str, str]] = None, query: Optional[Dict[str, str]] = None, + enriched_data: Optional[Dict[str, Any]] = None, ): - # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...). - harness.data["data"] = {"batch_id": batch_id} + # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...), + # then pre-call enrichment adds key/team metadata to it. + harness.data["data"] = {"batch_id": batch_id, **(enriched_data or {})} return await endpoints.retrieve_batch( request=FakeRequest(headers=headers, query=query), fastapi_response=Response(), @@ -1271,6 +1273,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness): "api_key": "sk-azure", "api_base": "https://azure.test", "model": "azure/gpt-4o", + "litellm_metadata": {"model_group": "azure/gpt-4o"}, } # 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip. @@ -1293,6 +1296,39 @@ async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment assert retrieve_harness.aretrieve_kwargs()["model"] == "azure/gpt-4o" +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__stamps_model_group(retrieve_harness): + """This path never goes through the router, so nothing else labels the call. + Without the stamp the spend log lands under a blank model group and the batch + disappears from per-model usage.""" + await call_retrieve(retrieve_harness, AZURE_BATCH_ID) + + litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"] + + assert litellm_metadata["model_group"] == "azure/gpt-4o" + + +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__stamps_model_group_beside_existing_metadata( + retrieve_harness, +): + """The stamp joins the metadata pre-call enrichment already built. Replacing + that dict instead of adding to it drops the key and team labels the spend log + is attributed with.""" + await call_retrieve( + retrieve_harness, + AZURE_BATCH_ID, + enriched_data={"litellm_metadata": {"user_api_key_alias": "team-a-key"}}, + ) + + litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"] + + assert litellm_metadata == { + "user_api_key_alias": "team-a-key", + "model_group": "azure/gpt-4o", + } + + @pytest.mark.asyncio async def test_retrieve__model_encoded_id__encodes_output_and_error_ids( retrieve_harness, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 836085c1f5d..e4e0dafd750 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1169,17 +1169,19 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") -async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]: +async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: loop = asyncio.get_event_loop() deadline = loop.time() + timeout while loop.time() < deadline: - keys = sorted( - key - for key in router.cache.in_memory_cache.cache_dict + cache_dict = router.cache.in_memory_cache.cache_dict + moved = sorted( + f"{key}={cache_dict[key]}" + for key in cache_dict if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) + and cache_dict[key] ) - if keys: - return keys + if moved: + return moved await asyncio.sleep(0.05) return [] @@ -1212,7 +1214,13 @@ def _batch_fan_out_router(routing_strategy: str): @pytest.mark.parametrize( "routing_strategy", - ["usage-based-routing", "latency-based-routing", "cost-based-routing", "least-busy"], + [ + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", + "least-busy", + ], ) @pytest.mark.asyncio async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( @@ -1239,10 +1247,10 @@ async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( for _ in range(3): response = await router.aretrieve_batch(batch_id=_BATCH_ID) await collector.retrieve_batch_payload() - strategy_keys = await _router_strategy_keys(router) + moved_counters = await _moved_routing_counters(router) assert response.id == _BATCH_ID - assert strategy_keys == [] + assert moved_counters == [] @pytest.mark.asyncio From 91761d984d55b2e4729e38c36ef7c7510c9ffc76 Mon Sep 17 00:00:00 2001 From: abhirup7 Date: Tue, 8 Sep 2026 00:30:00 +0530 Subject: [PATCH 013/349] fix(azure): send the resolved Entra ID token on image generation requests Azure image generation calls initialize_azure_sdk_client like the chat path does, but then sends the request through httpx with the headers it was given, so a credential resolved from litellm_params or the environment (Entra ID client credentials, managed or workload identity, OIDC, a static azure_ad_token) never reached the wire and Azure answered 401. Only an explicitly passed azure_ad_token_provider was applied Add get_azure_request_auth_headers, which turns the credential in azure_client_params into an Authorization: Bearer header (or api-key, following the SDK's precedence) while keeping any auth header the caller already set, and use it for both the sync and async image requests. The pre-call logging metadata receives a redacted copy of those headers so the token never reaches logging callbacks Fixes #16422 --- litellm/llms/azure/azure.py | 31 ++- litellm/llms/azure/common_utils.py | 35 +++ .../test_azure_image_generation_init.py | 249 ++++++++++++++++++ 3 files changed, 301 insertions(+), 14 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 46a9dd1a531..bfbaffd972c 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -46,7 +46,9 @@ from .common_utils import ( AzureOpenAIError, BaseAzureLLM, get_azure_ad_token_from_oidc, + get_azure_request_auth_headers, process_azure_headers, + redact_azure_auth_headers, select_azure_base_url_or_endpoint, ) from .image_generation import ( @@ -1142,7 +1144,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key: str, input: list, logging_obj: LiteLLMLoggingObj, - headers: dict, + headers: dict[str, str], client=None, timeout=None, model: str | None = None, @@ -1167,7 +1169,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={ "complete_input_dict": data, "api_base": img_gen_api_base, - "headers": headers, + "headers": redact_azure_auth_headers(headers), }, ) httpx_response: Final[httpx.Response] = await self.make_async_azure_httpx_request( @@ -1226,7 +1228,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): timeout: float, optional_params: dict, logging_obj: LiteLLMLoggingObj, - headers: dict, + headers: dict[str, str], model: str | None = None, api_key: str | None = None, api_base: str | None = None, @@ -1261,21 +1263,22 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if not isinstance(max_retries, int): raise AzureOpenAIError(status_code=422, message="max retries must be an int") - if api_key is None and azure_ad_token_provider is not None: - azure_ad_token = azure_ad_token_provider() - if azure_ad_token: - headers.pop("api-key", None) - headers["Authorization"] = f"Bearer {azure_ad_token}" - - # init AzureOpenAI Client + auth_params: Final[dict[str, object]] = {**(litellm_params or {})} # mutable-ok: SDK init takes a dict + if azure_ad_token is not None: + auth_params["azure_ad_token"] = azure_ad_token + if azure_ad_token_provider is not None: + auth_params["azure_ad_token_provider"] = azure_ad_token_provider azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client( - litellm_params=litellm_params or {}, + litellm_params=auth_params, api_key=api_key, model_name=model or "", api_version=api_version, api_base=api_base, is_async=False, ) + request_headers: Final = dict( # mutable-ok: the httpx request helpers take a dict + get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) + ) if aimg_generation is True: return self.aimage_generation( data=data, @@ -1286,7 +1289,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): client=client, azure_client_params=azure_client_params, timeout=timeout, - headers=headers, + headers=request_headers, model=model, ) @@ -1303,7 +1306,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={ "complete_input_dict": data, "api_base": img_gen_api_base, - "headers": headers, + "headers": redact_azure_auth_headers(request_headers), }, ) httpx_response: Final[httpx.Response] = self.make_sync_azure_httpx_request( @@ -1313,7 +1316,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_version=api_version or "", api_key=api_key or "", data=data, - headers=headers, + headers=request_headers, deployment_name=model, ) provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6cb7d09cec4..f276d8b18d1 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -405,6 +405,41 @@ def get_azure_ad_token( return azure_ad_token +_AZURE_AUTH_HEADER_NAMES: Final = frozenset(("api-key", "authorization")) +_REDACTED_AZURE_HEADER_VALUE: Final = "***REDACTED***" + + +def _resolve_azure_ad_token(azure_client_params: Mapping[str, object]) -> str | None: + azure_ad_token: Final = azure_client_params.get("azure_ad_token") + if isinstance(azure_ad_token, str) and azure_ad_token: + return azure_ad_token + token_provider: Final = azure_client_params.get("azure_ad_token_provider") + provided_token: Final = token_provider() if callable(token_provider) else None + return provided_token if isinstance(provided_token, str) and provided_token else None + + +def get_azure_request_auth_headers( + headers: Mapping[str, str], + azure_client_params: Mapping[str, object], +) -> Mapping[str, str]: + if any(name.lower() in _AZURE_AUTH_HEADER_NAMES for name in headers): + return headers + azure_ad_token: Final = _resolve_azure_ad_token(azure_client_params) + if azure_ad_token is not None: + return MappingProxyType({**headers, "Authorization": f"Bearer {azure_ad_token}"}) + api_key: Final = azure_client_params.get("api_key") + if isinstance(api_key, str) and api_key: + return MappingProxyType({**headers, "api-key": api_key}) + return headers + + +def redact_azure_auth_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return { # mutable-ok: logging callbacks JSON-serialize this copy + name: (_REDACTED_AZURE_HEADER_VALUE if name.lower() in _AZURE_AUTH_HEADER_NAMES else value) + for name, value in headers.items() + } + + class BaseAzureLLM(BaseOpenAILLM): @staticmethod def _try_get_default_azure_credential_provider( diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 70b5eab5c37..a30aa277f3d 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -10,6 +10,11 @@ import respx import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.azure.common_utils import ( + _cached_entra_id_token_provider, + get_azure_request_auth_headers, + redact_azure_auth_headers, +) from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, ) @@ -587,3 +592,247 @@ def test_azure_image_generation_v1_route_base_model_vs_deployment_name(respx_moc sent_body = json.loads(request.content) assert sent_body["model"] == model assert sent_body["prompt"] == prompt + + +@pytest.fixture +def fake_entra_id(monkeypatch: pytest.MonkeyPatch): + built_credentials = [] + + class FakeClientSecretCredential: + def __init__(self, tenant_id: str, client_id: str, client_secret: str) -> None: + built_credentials.append((tenant_id, client_id, client_secret)) + + monkeypatch.setattr("azure.identity.ClientSecretCredential", FakeClientSecretCredential) + monkeypatch.setattr("azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "entra-id-token") + _cached_entra_id_token_provider.cache_clear() + yield built_credentials + _cached_entra_id_token_provider.cache_clear() + + +def _mock_image_generation_route(respx_mock: respx.MockRouter, api_base: str, model: str) -> respx.Route: + return respx_mock.post(f"{api_base}/openai/deployments/{model}/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]}) + ) + + +@pytest.mark.parametrize("credentials_in_litellm_params", [False, True]) +def test_azure_image_generation_keyless_entra_id_sends_bearer_token( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fake_entra_id: list, + credentials_in_litellm_params: bool, +): + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + litellm_params = {"api_base": api_base, "api_version": api_version} + if credentials_in_litellm_params: + litellm_params.update( + tenant_id="tenant-from-params", client_id="client-from-params", client_secret="secret-from-params" + ) + expected_credential = ("tenant-from-params", "client-from-params", "secret-from-params") + else: + monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env") + expected_credential = ("tenant-from-env", "client-from-env", "secret-from-env") + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params=litellm_params, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-id-token" + assert "api-key" not in request.headers + assert fake_entra_id == [expected_credential] + assert response.data[0].b64_json == "aaaa" + logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"] + assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"} + assert "entra-id-token" not in str(logging_obj.pre_call.call_args) + + +@pytest.mark.asyncio +async def test_azure_aimage_generation_keyless_entra_id_sends_bearer_token( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = await AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + aimg_generation=True, + litellm_params={ + "api_base": api_base, + "api_version": api_version, + "tenant_id": "tenant-from-params", + "client_id": "client-from-params", + "client_secret": "secret-from-params", + }, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-id-token" + assert "api-key" not in request.headers + assert fake_entra_id == [("tenant-from-params", "client-from-params", "secret-from-params")] + assert response.data[0].b64_json == "aaaa" + logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"] + assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"} + assert "entra-id-token" not in str(logging_obj.pre_call.call_args) + + +@pytest.mark.parametrize( + "credential_kwargs, expected_authorization", + [ + ({"azure_ad_token": "static-ad-token"}, "Bearer static-ad-token"), + ({"azure_ad_token_provider": lambda: "provider-token"}, "Bearer provider-token"), + ], +) +def test_azure_image_generation_explicit_azure_ad_credential_sends_bearer_token( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + credential_kwargs: dict, + expected_authorization: str, +): + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + **credential_kwargs, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == expected_authorization + assert "api-key" not in request.headers + assert response.data[0].b64_json == "aaaa" + + +def test_azure_image_generation_with_api_key_keeps_api_key_header( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list +): + monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env") + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json", "api-key": "sk-test"}, + model="gpt-image-1", + api_key="sk-test", + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + ) + + request = route.calls.last.request + assert request.headers["api-key"] == "sk-test" + assert "Authorization" not in request.headers + assert fake_entra_id == [] + assert response.data[0].b64_json == "aaaa" + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***" + + +@pytest.mark.parametrize( + "caller_auth_header", + [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}], +) +def test_get_azure_request_auth_headers_keeps_caller_auth_header(caller_auth_header: dict): + headers = {"Content-Type": "application/json", **caller_auth_header} + azure_client_params = { + "api_key": "sk-resolved", + "azure_ad_token": "resolved-token", + "azure_ad_token_provider": lambda: "provider-token", + } + assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers + + +def test_get_azure_request_auth_headers_prefers_azure_ad_token_over_provider_and_api_key(): + headers = {"Content-Type": "application/json"} + azure_client_params = { + "api_key": "sk-resolved", + "azure_ad_token": "static-token", + "azure_ad_token_provider": lambda: "provider-token", + } + out = get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) + assert dict(out) == {"Content-Type": "application/json", "Authorization": "Bearer static-token"} + assert headers == {"Content-Type": "application/json"} + + +def test_get_azure_request_auth_headers_uses_token_provider_over_api_key(): + azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": lambda: "pt"} + out = get_azure_request_auth_headers(headers={}, azure_client_params=azure_client_params) + assert dict(out) == {"Authorization": "Bearer pt"} + + +def test_get_azure_request_auth_headers_falls_back_to_api_key(): + azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": None} + out = get_azure_request_auth_headers(headers={"Content-Type": "application/json"}, azure_client_params=azure_client_params) + assert dict(out) == {"Content-Type": "application/json", "api-key": "sk-resolved"} + + +@pytest.mark.parametrize( + "azure_client_params", + [ + {}, + {"api_key": "", "azure_ad_token": "", "azure_ad_token_provider": None}, + {"azure_ad_token_provider": lambda: None}, + {"azure_ad_token_provider": lambda: ""}, + ], +) +def test_get_azure_request_auth_headers_without_credential_leaves_headers_unchanged(azure_client_params: dict): + headers = {"Content-Type": "application/json"} + assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers + + +def test_redact_azure_auth_headers_masks_only_credential_values(): + headers = {"Content-Type": "application/json", "api-key": "sk-secret", "authorization": "Bearer secret"} + assert redact_azure_auth_headers(headers) == { + "Content-Type": "application/json", + "api-key": "***REDACTED***", + "authorization": "***REDACTED***", + } + assert headers["api-key"] == "sk-secret" + assert headers["authorization"] == "Bearer secret" From 66ebc722d6751a7a8d0aef751e8d2e5f6a2efe49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:33:09 -0700 Subject: [PATCH 014/349] perf(azure): reuse the token refresh credential across image generation requests With enable_azure_ad_token_refresh, every keyless image request built a new DefaultAzureCredential and fetched a token. Cache the provider per scope like the Entra ID one. --- litellm/llms/azure/common_utils.py | 9 ++-- .../test_azure_image_generation_init.py | 47 +++++++++++++++++++ .../llms/azure/test_azure_common_utils.py | 3 ++ 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index f276d8b18d1..9f70761514b 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -95,6 +95,11 @@ def _cached_entra_id_token_provider( return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) +@lru_cache(maxsize=128) +def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=scope) + + def get_azure_ad_token_from_entra_id( tenant_id: str, client_id: str, @@ -649,9 +654,7 @@ class BaseAzureLLM(BaseOpenAILLM): "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth" ) try: - azure_ad_token_provider = get_azure_ad_token_provider( - azure_scope=scope, - ) + azure_ad_token_provider = _cached_azure_ad_token_refresh_provider(scope) except ValueError: verbose_logger.debug("Azure AD Token Provider could not be used.") if api_version is None: diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index a30aa277f3d..cfde1760389 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -11,6 +11,7 @@ import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.common_utils import ( + _cached_azure_ad_token_refresh_provider, _cached_entra_id_token_provider, get_azure_request_auth_headers, redact_azure_auth_headers, @@ -775,6 +776,52 @@ def test_azure_image_generation_with_api_key_keeps_api_key_header( assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***" +@pytest.fixture +def fake_default_azure_credential(monkeypatch: pytest.MonkeyPatch): + built_credentials = [] + + class FakeDefaultAzureCredential: + def __init__(self) -> None: + built_credentials.append(self) + + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", "AZURE_AD_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", FakeDefaultAzureCredential) + monkeypatch.setattr( + "azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "default-credential-token" + ) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + _cached_azure_ad_token_refresh_provider.cache_clear() + yield built_credentials + _cached_azure_ad_token_refresh_provider.cache_clear() + + +def test_azure_image_generation_token_refresh_reuses_credential_across_requests( + respx_mock: respx.MockRouter, fake_default_azure_credential: list +): + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + + for _ in range(3): + AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + ) + + assert route.call_count == 3 + assert all(call.request.headers["Authorization"] == "Bearer default-credential-token" for call in route.calls) + assert len(fake_default_azure_credential) == 1 + + @pytest.mark.parametrize( "caller_auth_header", [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}], diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f000abb4c9a..7189be7c052 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, + _cached_azure_ad_token_refresh_provider, _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, @@ -34,6 +35,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_TENANT_ID", raising=False) monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) + _cached_azure_ad_token_refresh_provider.cache_clear() with ( patch( @@ -78,6 +80,7 @@ def setup_mocks(monkeypatch): "logger": mock_logger, "select_url": mock_select_url, } + _cached_azure_ad_token_refresh_provider.cache_clear() def test_initialize_with_api_key(setup_mocks): From 233337628f4b6f1c9ec0527d5d442cfe512ac08b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:06:25 -0400 Subject: [PATCH 015/349] feat(batches): support Mistral files/batches and per-page OCR batch cost tracking Adds MistralFilesConfig and MistralBatchesConfig so Mistral can be used as a Files and Batches provider through the shared BaseLLMHTTPHandler path, the same way Bedrock plugs in. /v1/ocr is now an accepted batch endpoint, and completed OCR batches are billed per page (ocr_cost_per_page_batches, half the synchronous rate) instead of per token. Resolves #29914 --- litellm/batches/batch_utils.py | 41 ++- litellm/batches/main.py | 20 +- litellm/cost_calculator.py | 61 ++++ litellm/files/main.py | 9 +- litellm/files/types.py | 2 +- litellm/llms/mistral/batches/__init__.py | 0 .../llms/mistral/batches/transformation.py | 186 +++++++++++++ litellm/llms/mistral/common_utils.py | 36 +++ litellm/llms/mistral/files/__init__.py | 0 litellm/llms/mistral/files/transformation.py | 226 +++++++++++++++ ...odel_prices_and_context_window_backup.json | 40 ++- litellm/types/llms/openai.py | 2 +- litellm/types/utils.py | 4 + litellm/utils.py | 10 + model_prices_and_context_window.json | 40 ++- .../test_litellm/batches/test_batch_utils.py | 83 ++++++ tests/test_litellm/batches/test_main.py | 42 +++ .../llms/mistral/batches/__init__.py | 0 .../test_mistral_batches_transformation.py | 260 ++++++++++++++++++ .../llms/mistral/files/__init__.py | 0 .../test_mistral_files_transformation.py | 189 +++++++++++++ .../llms/mistral/ocr/test_mistral_ocr_cost.py | 4 +- tests/test_litellm/test_utils.py | 2 + 23 files changed, 1216 insertions(+), 41 deletions(-) create mode 100644 litellm/llms/mistral/batches/__init__.py create mode 100644 litellm/llms/mistral/batches/transformation.py create mode 100644 litellm/llms/mistral/common_utils.py create mode 100644 litellm/llms/mistral/files/__init__.py create mode 100644 litellm/llms/mistral/files/transformation.py create mode 100644 tests/test_litellm/llms/mistral/batches/__init__.py create mode 100644 tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py create mode 100644 tests/test_litellm/llms/mistral/files/__init__.py create mode 100644 tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 87f8fd3946e..077d6e72fd7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details +from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -50,7 +51,7 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -80,7 +81,7 @@ async def calculate_batch_cost_and_usage( async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, litellm_params: dict | None = None, model_info: ModelInfo | None = None, @@ -166,7 +167,7 @@ class _BatchOutputLineStats: def _classify_output_line_stats( entries: Iterable[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats | _LineOutcome]: @@ -185,7 +186,7 @@ def _classify_output_line_stats( def _safe_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats | None: @@ -207,7 +208,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats: @@ -218,6 +219,7 @@ def _compute_output_line_stats( response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details line_prompt_cost, line_completion_cost = _output_line_cost( + response_body=response_body, usage=usage, custom_llm_provider=custom_llm_provider, model_name=model_name, @@ -237,19 +239,36 @@ def _compute_output_line_stats( ) +def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> OCRUsageInfo | None: + """OCR results report ``usage_info`` (pages) instead of ``usage`` (tokens); None for non-OCR lines.""" + raw_usage_info: Final = response_body.get("usage_info") + if not isinstance(raw_usage_info, Mapping): + return None + return OCRUsageInfo.model_validate(raw_usage_info) + + def _output_line_cost( + response_body: Mapping[str, object], usage: Usage, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, ) -> tuple[float, float]: """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" - from litellm.cost_calculator import batch_cost_calculator + from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost cost_model: Final = ( model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" ) + ocr_usage: Final = _ocr_usage_info_from_response_body(response_body) + if ocr_usage is not None: + return ocr_batch_cost( + model=cost_model, + custom_llm_provider=custom_llm_provider, + usage_info=ocr_usage, + model_info=model_info, + ) return batch_cost_calculator( usage=usage, model=cost_model, @@ -260,7 +279,7 @@ def _output_line_cost( def _aggregate_batch_cost_usage_models( entries: Iterable[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -427,7 +446,7 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -457,7 +476,7 @@ async def _fetch_batch_managed_file_content( async def _fetch_batch_output_file_content( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -479,7 +498,7 @@ async def _fetch_batch_output_file_content( async def count_error_file_failed_requests( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], litellm_params: dict | None, ) -> int: """Count failed requests reported only in the batch's separate error file. diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 77a4fdebf16..76b6c73b375 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -105,9 +105,11 @@ def _resolve_timeout( @client async def acreate_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -155,9 +157,11 @@ async def acreate_batch( @client def create_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -341,7 +345,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -389,7 +393,7 @@ def _handle_retrieve_batch_providers_without_provider_config( _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", logging_obj: LiteLLMLoggingObj | None = None, ): @@ -497,7 +501,7 @@ def _handle_retrieve_batch_providers_without_provider_config( message=( f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. " "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. " - "'bedrock' is supported but requires `model` to be passed so the provider config can be loaded." + "'bedrock' and 'mistral' are supported but require `model` to be passed so the provider config can be loaded." ), model="n/a", llm_provider=custom_llm_provider, @@ -514,7 +518,7 @@ def _handle_retrieve_batch_providers_without_provider_config( def retrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..0c0c6b05df8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -139,6 +139,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LitellmLoggingObject, ) + from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo else: LitellmLoggingObject = Any @@ -1982,6 +1983,66 @@ def ocr_cost( return ocr_pages_cost + annotation_pages_cost, 0.0 +_OCR_PRICING_KEYS: Final = ( + "ocr_cost_per_page", + "ocr_cost_per_page_batches", + "annotation_cost_per_page", + "annotation_cost_per_page_batches", +) + + +def ocr_batch_cost( + model: str, + custom_llm_provider: str | None, + usage_info: "OCRUsageInfo", + model_info: ModelInfo | None = None, +) -> tuple[float, float]: + """Per-page cost of one OCR result inside a batch output file. + + Batch OCR is billed per page at the ``*_batches`` rate, falling back to the + synchronous per-page rate when a model has no batch price recorded, the same + fallback ``batch_cost_calculator`` applies to per-token batch pricing. Returns + ``(prompt_cost, completion_cost)`` with the whole cost in the first slot, like + ``ocr_cost``. + """ + has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) + if has_ocr_pricing: + resolved_info: ModelInfo | None = model_info + else: + try: + resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + resolved_info = None + if resolved_info is None: + verbose_logger.warning( + "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", + model, + custom_llm_provider, + ) + return 0.0, 0.0 + + page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") + annotation_rate: Final = _first_price( + resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page" + ) + pages_processed: Final = usage_info.pages_processed or 0 + annotation_pages: Final = usage_info.pages_processed_annotation or 0 + if page_rate is None and pages_processed > 0: + verbose_logger.warning( + "OCR batch cost: model=%s custom_llm_provider=%s reported pages_processed=%s but no " + "ocr_cost_per_page is configured; returning 0.0 cost for those pages.", + model, + custom_llm_provider, + pages_processed, + ) + effective_annotation_rate: Final = annotation_rate if annotation_rate is not None else page_rate + return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0 + + +def _first_price(model_info: ModelInfo, *keys: str) -> float | None: + return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) + + def vector_store_search_cost( model: str | None, custom_llm_provider: str, diff --git a/litellm/files/main.py b/litellm/files/main.py index 218518eb3cd..3d90bf4f299 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -27,12 +27,15 @@ FileCreateProvider = Literal[ "litellm_proxy", "manus", "anthropic", + "mistral", ] FileRetrieveProvider = Literal[ - "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"] -FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] +FileDeleteProvider = Literal[ + "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral" +] +FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse diff --git a/litellm/files/types.py b/litellm/files/types.py index b4ec9996f37..01c7970144b 100644 --- a/litellm/files/types.py +++ b/litellm/files/types.py @@ -2,7 +2,7 @@ from collections.abc import AsyncIterator, Iterator from typing import Literal, NamedTuple FileContentProvider = Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus", "mistral" ] diff --git a/litellm/llms/mistral/batches/__init__.py b/litellm/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py new file mode 100644 index 00000000000..399319e590b --- /dev/null +++ b/litellm/llms/mistral/batches/transformation.py @@ -0,0 +1,186 @@ +""" +Mistral Batch API. Reference: https://docs.mistral.ai/api/#tag/batch + +Mistral runs one model per job (set on the job, not per input line) and accepts +``/v1/ocr`` as a batch endpoint, which is how OCR gets its 50% batch discount. +Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_code, body}}``), +so the shared batch cost accounting reads them without a provider branch. +""" + +from types import MappingProxyType +from typing import Final, Literal + +import httpx +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Errors as BatchErrors +from openai.types.batch_error import BatchError +from pydantic import BaseModel, ConfigDict + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders + +from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error + +MistralBatchStatus = Literal[ + "QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED" +] +OpenAIBatchStatus = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + +_STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType( + { + "QUEUED": "validating", + "RUNNING": "in_progress", + "SUCCESS": "completed", + "FAILED": "failed", + "TIMEOUT_EXCEEDED": "expired", + "CANCELLATION_REQUESTED": "cancelling", + "CANCELLED": "cancelled", + } +) + + +class MistralBatchError(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + message: str + count: int = 1 + + +class MistralBatchJob(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + input_files: tuple[str, ...] = () + endpoint: str + model: str | None = None + status: MistralBatchStatus + created_at: int + started_at: int | None = None + completed_at: int | None = None + total_requests: int = 0 + completed_requests: int = 0 + succeeded_requests: int = 0 + failed_requests: int = 0 + output_file: str | None = None + error_file: str | None = None + errors: tuple[MistralBatchError, ...] = () + metadata: dict[str, str] | None = None + + +def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: + status: Final = _STATUS_MAP[job.status] + terminal_at: Final = job.completed_at + return LiteLLMBatch( + id=job.id, + object="batch", + endpoint=job.endpoint, + input_file_id=job.input_files[0] if job.input_files else "", + completion_window="24h", + status=status, + created_at=job.created_at, + in_progress_at=job.started_at, + completed_at=terminal_at if status == "completed" else None, + failed_at=terminal_at if status == "failed" else None, + expired_at=terminal_at if status == "expired" else None, + cancelled_at=terminal_at if status == "cancelled" else None, + output_file_id=job.output_file, + error_file_id=job.error_file, + errors=( + BatchErrors( + object="list", + data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors], + ) + if job.errors + else None + ), + request_counts=BatchRequestCounts( + total=job.total_requests, + completed=job.succeeded_requests, + failed=job.failed_requests, + ), + metadata=job.metadata, + ) + + +class MistralBatchesConfig(BaseBatchesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.MISTRAL + + def validate_environment( + self, + headers: dict, + model: str, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + return get_mistral_auth_headers(headers, api_key) + + def get_complete_batch_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + data: CreateBatchRequest, + ) -> str: + return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" + + def transform_create_batch_request( + self, + model: str, + create_batch_data: CreateBatchRequest, + optional_params: dict, + litellm_params: dict, + ) -> dict[str, object]: + metadata: Final = create_batch_data.get("metadata") + return { + "input_files": [create_batch_data["input_file_id"]], + "endpoint": create_batch_data["endpoint"], + "model": model, + **({"metadata": metadata} if metadata else {}), + **(create_batch_data.get("extra_body") or {}), + } + + def transform_create_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: dict, + ) -> LiteLLMBatch: + return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) + + def transform_retrieve_batch_request( + self, + batch_id: str, + optional_params: dict, + litellm_params: dict, + ) -> dict[str, object]: + encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") + return { + "method": "GET", + "url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}", + "headers": get_mistral_auth_headers({}, litellm_params.get("api_key")), + } + + def transform_retrieve_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: dict, + ) -> LiteLLMBatch: + return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) + + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + return mistral_error(error_message, status_code, headers) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py new file mode 100644 index 00000000000..9ea501c860d --- /dev/null +++ b/litellm/llms/mistral/common_utils.py @@ -0,0 +1,36 @@ +from typing import Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str + +MISTRAL_API_BASE: Final = "https://api.mistral.ai" +MISTRAL_API_KEY_ENV_VAR: Final = "MISTRAL_API_KEY" + + +class MistralError(BaseLLMException): + pass + + +def get_mistral_api_base(api_base: str | None) -> str: + """Return the Mistral origin without a trailing ``/v1``, so callers can append ``/v1/``.""" + resolved: Final = (api_base or get_secret_str("MISTRAL_API_BASE") or MISTRAL_API_BASE).rstrip("/") + return resolved.removesuffix("/v1") + + +def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict: + resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR) + if resolved_key is None: + raise ValueError( + "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" + ) + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + +def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError: + return MistralError( + status_code=status_code, + message=error_message, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), + ) diff --git a/litellm/llms/mistral/files/__init__.py b/litellm/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py new file mode 100644 index 00000000000..071b6f58569 --- /dev/null +++ b/litellm/llms/mistral/files/transformation.py @@ -0,0 +1,226 @@ +""" +Mistral Files API. Reference: https://docs.mistral.ai/api/#tag/files + +Mistral's file objects already carry the OpenAI field names (id, bytes, created_at, +filename, purpose), so this config is URL routing, auth, and a purpose mapping: +Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes. +""" + +import time +from typing import Final, Literal + +import httpx +from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, ConfigDict + +from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj +from litellm.types.llms.openai import ( + CreateFileRequest, + FileContentRequest, + HttpxBinaryResponseContent, + OpenAICreateFileRequestOptionalParams, + OpenAIFileObject, + OpenAIFilesPurpose, +) +from litellm.types.utils import LlmProviders + +from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error + +MistralFilePurpose = Literal["fine-tune", "batch", "ocr"] + + +class MistralFile(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + bytes: int = 0 + created_at: int | None = None + filename: str = "" + purpose: MistralFilePurpose = "batch" + expires_at: int | None = None + + +class MistralFileList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[MistralFile, ...] = () + + +class MistralFileDeleted(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + deleted: bool = True + + +def _to_openai_file_object(file: MistralFile) -> OpenAIFileObject: + return OpenAIFileObject( + id=file.id, + bytes=file.bytes, + created_at=file.created_at if file.created_at is not None else int(time.time()), + filename=file.filename, + object="file", + purpose=_to_openai_purpose(file.purpose), + status="uploaded", + expires_at=file.expires_at, + ) + + +def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose: + match purpose: + case "fine-tune" | "batch": + return purpose + case "ocr": + return "user_data" + + +def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: + match purpose: + case "fine-tune" | "ocr": + return purpose + case _: + return "batch" + + +class MistralFilesConfig(BaseFilesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.MISTRAL + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + stream: bool | None = None, + ) -> str: + return f"{get_mistral_api_base(api_base)}/v1/files" + + def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str: + encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") + return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}" + + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + return mistral_error(error_message, status_code, headers) + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + return get_mistral_auth_headers(headers, api_key) + + def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: + return ["purpose"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + return optional_params + + def transform_create_file_request( + self, + model: str, + create_file_data: CreateFileRequest, + optional_params: dict, + litellm_params: dict, + ) -> dict: + file_data: Final = create_file_data.get("file") + if file_data is None: + raise ValueError("File data is required") + extracted: Final = extract_file_data(file_data) + filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" + content_type: Final = extracted.get("content_type") or "application/octet-stream" + return { + "file": (filename, extracted["content"], content_type), + "purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + } + + def transform_create_file_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) + + def transform_retrieve_file_request( + self, + file_id: str, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_id, litellm_params), {} + + def transform_retrieve_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) + + def transform_delete_file_request( + self, + file_id: str, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_id, litellm_params), {} + + def transform_delete_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> FileDeleted: + deleted: Final = MistralFileDeleted.model_validate(raw_response.json()) + return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") + + def transform_list_files_request( + self, + purpose: str | None, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {} + return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params + + def transform_list_files_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> list[OpenAIFileObject]: + return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data] + + def transform_file_content_request( + self, + file_content_request: FileContentRequest, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {} + + def transform_file_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 54ebdc85be9..5a934301edd 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -35262,51 +35262,66 @@ "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-1": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, + "ocr_cost_per_page_batches": 0.0005, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-2512": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, @@ -59822,31 +59837,40 @@ "mistral/mistral-ocr-3": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-3-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/voxtral-mini-latest": { diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..defd59f2be8 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -498,7 +498,7 @@ class CreateBatchRequest(TypedDict, total=False): """ completion_window: Literal["24h"] - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"] + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"] input_file_id: str metadata: dict[str, str] | None output_expires_after: FileExpiresAfter diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d62f00f3676..ea23e00d2bb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -320,8 +320,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models + ocr_cost_per_page_batches: ReadOnly[float | None] ocr_cost_per_credit: float | None # for OCR models priced by credit annotation_cost_per_page: float | None # for OCR models + annotation_cost_per_page_batches: ReadOnly[float | None] search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None @@ -3598,8 +3600,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_512k_tokens: float | None = None output_vector_size: int | None = None ocr_cost_per_page: float | None = None + ocr_cost_per_page_batches: float | None = None ocr_cost_per_credit: float | None = None annotation_cost_per_page: float | None = None + annotation_cost_per_page_batches: float | None = None regional_processing_uplift_multiplier_eu: float | None = None regional_processing_uplift_multiplier_us: float | None = None regional_endpoint_uplift_multiplier: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 8df28870544..3ca81605a95 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5963,8 +5963,10 @@ def _get_model_info_helper( tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None), + ocr_cost_per_page_batches=_model_info.get("ocr_cost_per_page_batches", None), ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None), annotation_cost_per_page=_model_info.get("annotation_cost_per_page", None), + annotation_cost_per_page_batches=_model_info.get("annotation_cost_per_page_batches", None), provider_specific_entry=_model_info.get("provider_specific_entry", None), uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), @@ -8909,6 +8911,10 @@ class ProviderConfigManager: from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig return AnthropicFilesConfig() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.files.transformation import MistralFilesConfig + + return MistralFilesConfig() return None @staticmethod @@ -8920,6 +8926,10 @@ class ProviderConfigManager: from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig return BedrockBatchesConfig() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.batches.transformation import MistralBatchesConfig + + return MistralBatchesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 54ebdc85be9..5a934301edd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -35262,51 +35262,66 @@ "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-1": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, + "ocr_cost_per_page_batches": 0.0005, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-2512": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, @@ -59822,31 +59837,40 @@ "mistral/mistral-ocr-3": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-3-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/voxtral-mini-latest": { diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 768ea332677..85f267c2db0 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1787,3 +1787,86 @@ class TestBatchCostIsFinal: @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) def test_other_terminal_statuses_are_final(self, status): assert bu.batch_cost_is_final(_retrieved_batch(status)) is True + + +# =========================================================================== # +# OCR batch output lines (Mistral /v1/ocr batches) are billed per page, not per token +# =========================================================================== # + + +def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): + usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} + if annotation_pages is not None: + usage_info["pages_processed_annotation"] = annotation_pages + return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info) + + +def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002}, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(3), _ocr_row(5), _failed_row(model="mistral-ocr-latest")], + custom_llm_provider="mistral", + model_name="mistral/mistral-ocr-latest", + ) + assert result.cost == pytest.approx(8 * 0.002) + assert result.prompt_cost == pytest.approx(8 * 0.002) + assert result.completion_cost == 0.0 + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert result.usage.total_tokens == 0 + assert result.models == ["mistral/mistral-ocr-latest"] + + +def test_ocr_rows_fall_back_to_sync_page_rate_without_batch_price(monkeypatch): + monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004}) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(2)], custom_llm_provider="mistral") + assert result.cost == pytest.approx(2 * 0.004) + + +def test_ocr_rows_bill_annotation_pages_separately(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: { + "ocr_cost_per_page_batches": 0.002, + "annotation_cost_per_page_batches": 0.0025, + }, + ) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral") + assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025) + + +def test_ocr_rows_use_deployment_model_info_pricing_over_cost_map(monkeypatch): + monkeypatch.setattr( + litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted") + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(10)], + custom_llm_provider="mistral", + model_info={"ocr_cost_per_page_batches": 0.001}, + ) + assert result.cost == pytest.approx(0.01) + + +def test_ocr_rows_without_pricing_bill_zero_but_count_as_successful(monkeypatch): + monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"mode": "ocr"}) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(3)], custom_llm_provider="mistral") + assert result.cost == 0.0 + assert (result.successful_requests, result.failed_requests) == (1, 0) + + +def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_success_row(model="mistral-small-latest", usage=_usage(10, 5))], + custom_llm_provider="mistral", + ) + assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2) + assert result.usage.total_tokens == 15 diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index b87f9489250..c5a33dd6508 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -778,3 +778,45 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] assert "_litellm_internal_model_credentials" not in litellm_params + + +# =========================================================================== # +# mistral - a provider-config provider, like bedrock, so it requires `model` +# =========================================================================== # + + +def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams): + with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg: + result = bm.create_batch( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-abc", + custom_llm_provider="mistral", + model="mistral/mistral-ocr-latest", + ) + + assert result is seams.base_http.create_batch.return_value + _assert_only(seams.base_http.create_batch, seams, "create_batch") + get_cfg.assert_called_once() + forwarded = seams.base_http.create_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" + assert forwarded["model"] == "mistral-ocr-latest" + assert forwarded["create_batch_data"]["endpoint"] == "/v1/ocr" + + +def test_create__mistral_without_model_raises_badrequest(seams): + with pytest.raises(litellm.exceptions.BadRequestError): + bm.create_batch(**CREATE_KW, custom_llm_provider="mistral") + + for m in _all_seam_methods(seams, "create_batch"): + m.assert_not_called() + + +def test_retrieve__mistral_routes_to_base_http_handler_with_mistral_config(seams): + result = bm.retrieve_batch(batch_id="job-1", custom_llm_provider="mistral", model="mistral/mistral-ocr-latest") + + assert result is seams.base_http.retrieve_batch.return_value + _assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch") + forwarded = seams.base_http.retrieve_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" + assert forwarded["batch_id"] == "job-1" diff --git a/tests/test_litellm/llms/mistral/batches/__init__.py b/tests/test_litellm/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py new file mode 100644 index 00000000000..03e9c351a30 --- /dev/null +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -0,0 +1,260 @@ +""" +Regression tests for ``MistralBatchesConfig``, the BaseBatchesConfig implementation +behind ``custom_llm_provider="mistral"`` on /v1/batches. + +Locks the request shape Mistral's ``POST /v1/batch/jobs`` accepts (input_files list, +model set on the job, endpoint passed through untouched so ``/v1/ocr`` batches work), +the Mistral -> OpenAI status mapping, request-count and file-id mapping, and auth. +Everything runs for real against canned httpx responses; only the API key env var is +set. +""" + +import json + +import httpx +import pytest + +from litellm.llms.mistral.batches.transformation import MistralBatchesConfig +from litellm.llms.mistral.common_utils import MistralError +from litellm.types.llms.openai import CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders + +STATUS_MAP = { + "QUEUED": "validating", + "RUNNING": "in_progress", + "SUCCESS": "completed", + "FAILED": "failed", + "TIMEOUT_EXCEEDED": "expired", + "CANCELLATION_REQUESTED": "cancelling", + "CANCELLED": "cancelled", +} + + +def _job(**overrides): + base = { + "id": "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b", + "object": "batch", + "input_files": ["c1a2b3d4-0000-4000-8000-000000000001"], + "endpoint": "/v1/ocr", + "model": "mistral-ocr-latest", + "status": "SUCCESS", + "created_at": 1_757_400_000, + "started_at": 1_757_400_010, + "completed_at": 1_757_400_500, + "total_requests": 3, + "completed_requests": 3, + "succeeded_requests": 2, + "failed_requests": 1, + "output_file": "out-0000-4000-8000-000000000002", + "error_file": "err-0000-4000-8000-000000000003", + "errors": [], + "metadata": {"job_type": "testing"}, + } + return {**base, **overrides} + + +def _response(payload: dict, status_code: int = 200) -> httpx.Response: + return httpx.Response( + status_code=status_code, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://api.mistral.ai/v1/batch/jobs/x"), + ) + + +@pytest.fixture +def config() -> MistralBatchesConfig: + return MistralBatchesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + return "sk-mistral-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.MISTRAL + + +# --------------------------------------------------------------------------- # +# create +# --------------------------------------------------------------------------- # + + +def test_create_request_maps_openai_fields_onto_mistral_job(config): + data = CreateBatchRequest( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-123", + metadata={"team": "docs"}, + ) + body = config.transform_create_batch_request( + model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={} + ) + assert body == { + "input_files": ["file-123"], + "endpoint": "/v1/ocr", + "model": "mistral-ocr-latest", + "metadata": {"team": "docs"}, + } + + +def test_create_request_omits_empty_metadata_and_forwards_extra_body(config): + data = CreateBatchRequest( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-123", + metadata=None, + extra_body={"timeout_hours": 48}, + ) + body = config.transform_create_batch_request( + model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={} + ) + assert "metadata" not in body + assert body["timeout_hours"] == 48 + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.mistral.ai/v1/batch/jobs"), + ("https://api.mistral.ai/v1", "https://api.mistral.ai/v1/batch/jobs"), + ("https://proxy.example.com/", "https://proxy.example.com/v1/batch/jobs"), + ], +) +def test_create_url(config, api_base, expected): + url = config.get_complete_batch_url( + api_base=api_base, api_key="k", model="m", optional_params={}, litellm_params={}, data={} + ) + assert url == expected + + +def test_validate_environment_uses_bearer_auth(config, api_key): + headers = config.validate_environment( + headers={"x-extra": "1"}, model="m", messages=[], optional_params={}, litellm_params={} + ) + assert headers == {"x-extra": "1", "Authorization": f"Bearer {api_key}"} + + +def test_validate_environment_explicit_key_wins(config, api_key): + headers = config.validate_environment( + headers={}, model="m", messages=[], optional_params={}, litellm_params={}, api_key="sk-explicit" + ) + assert headers["Authorization"] == "Bearer sk-explicit" + + +def test_validate_environment_without_key_raises(config, monkeypatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + with pytest.raises(ValueError, match="Missing Mistral API Key"): + config.validate_environment(headers={}, model="m", messages=[], optional_params={}, litellm_params={}) + + +def test_create_response_maps_job_onto_openai_batch(config): + batch = config.transform_create_batch_response( + model="mistral-ocr-latest", + raw_response=_response(_job(status="QUEUED", started_at=None, completed_at=None)), + logging_obj=None, + litellm_params={}, + ) + assert isinstance(batch, LiteLLMBatch) + assert batch.id == "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b" + assert batch.endpoint == "/v1/ocr" + assert batch.input_file_id == "c1a2b3d4-0000-4000-8000-000000000001" + assert batch.status == "validating" + assert batch.created_at == 1_757_400_000 + assert batch.in_progress_at is None + assert batch.completed_at is None + assert batch.metadata == {"job_type": "testing"} + + +# --------------------------------------------------------------------------- # +# retrieve +# --------------------------------------------------------------------------- # + + +def test_retrieve_request_is_presigned_get_with_auth(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"} + ) + assert req["method"] == "GET" + assert req["url"] == "https://api.mistral.ai/v1/batch/jobs/job%2Fwith%20slash" + assert req["headers"] == {"Authorization": f"Bearer {api_key}"} + + +def test_retrieve_request_prefers_litellm_params_api_key(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="job-1", optional_params={}, litellm_params={"api_key": "sk-from-deployment"} + ) + assert req["headers"]["Authorization"] == "Bearer sk-from-deployment" + + +@pytest.mark.parametrize("mistral_status,openai_status", sorted(STATUS_MAP.items())) +def test_retrieve_response_status_mapping(config, mistral_status, openai_status): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={} + ) + assert batch.status == openai_status + + +@pytest.mark.parametrize( + "mistral_status,populated_field", + [ + ("SUCCESS", "completed_at"), + ("FAILED", "failed_at"), + ("TIMEOUT_EXCEEDED", "expired_at"), + ("CANCELLED", "cancelled_at"), + ], +) +def test_retrieve_response_terminal_timestamp_lands_on_matching_field(config, mistral_status, populated_field): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={} + ) + terminal_fields = {"completed_at", "failed_at", "expired_at", "cancelled_at"} + assert getattr(batch, populated_field) == 1_757_400_500 + for other in terminal_fields - {populated_field}: + assert getattr(batch, other) is None + assert batch.in_progress_at == 1_757_400_010 + + +def test_retrieve_response_maps_counts_and_files(config): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job()), logging_obj=None, litellm_params={} + ) + assert batch.request_counts.total == 3 + assert batch.request_counts.completed == 2 + assert batch.request_counts.failed == 1 + assert batch.output_file_id == "out-0000-4000-8000-000000000002" + assert batch.error_file_id == "err-0000-4000-8000-000000000003" + assert batch.errors is None + + +def test_retrieve_response_surfaces_job_errors(config): + batch = config.transform_retrieve_batch_response( + model=None, + raw_response=_response( + _job(status="FAILED", errors=[{"message": "invalid document", "count": 2}, {"message": "timeout"}]) + ), + logging_obj=None, + litellm_params={}, + ) + assert [e.message for e in batch.errors.data] == ["invalid document (x2)", "timeout"] + + +def test_retrieve_response_without_files_or_input(config): + batch = config.transform_retrieve_batch_response( + model=None, + raw_response=_response(_job(input_files=[], output_file=None, error_file=None, metadata=None)), + logging_obj=None, + litellm_params={}, + ) + assert batch.input_file_id == "" + assert batch.output_file_id is None + assert batch.error_file_id is None + assert batch.metadata is None + + +def test_get_error_class(config): + err = config.get_error_class("nope", 401, {"x-request-id": "r1"}) + assert isinstance(err, MistralError) + assert err.status_code == 401 + assert err.message == "nope" diff --git a/tests/test_litellm/llms/mistral/files/__init__.py b/tests/test_litellm/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py new file mode 100644 index 00000000000..d6ad8a34b35 --- /dev/null +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -0,0 +1,189 @@ +""" +Regression tests for ``MistralFilesConfig``, the BaseFilesConfig implementation behind +``custom_llm_provider="mistral"`` on /v1/files. + +Locks the URL routing for each file operation, the multipart upload shape Mistral's +``POST /v1/files`` accepts (purpose restricted to fine-tune/batch/ocr), and the +Mistral -> OpenAI file object mapping. Runs against canned httpx responses. +""" + +import json + +import httpx +import pytest +from openai.types.file_deleted import FileDeleted + +from litellm.llms.mistral.files.transformation import MistralFilesConfig +from litellm.types.llms.openai import CreateFileRequest, FileContentRequest, OpenAIFileObject +from litellm.types.utils import LlmProviders + +FILE_ID = "497f6eca-6276-4993-bfeb-53cbbbba6f09" + + +def _file(**overrides): + base = { + "id": FILE_ID, + "object": "file", + "bytes": 13000, + "created_at": 1_716_963_433, + "filename": "batch_input.jsonl", + "purpose": "batch", + "sample_type": "batch_request", + "num_lines": 3, + "source": "upload", + } + return {**base, **overrides} + + +def _response(payload) -> httpx.Response: + return httpx.Response( + status_code=200, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://api.mistral.ai/v1/files"), + ) + + +@pytest.fixture +def config() -> MistralFilesConfig: + return MistralFilesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + return "sk-mistral-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.MISTRAL + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.mistral.ai/v1/files"), + ("https://api.mistral.ai/v1/", "https://api.mistral.ai/v1/files"), + ("https://proxy.example.com", "https://proxy.example.com/v1/files"), + ], +) +def test_upload_url(config, api_base, expected): + url = config.get_complete_url(api_base=api_base, api_key="k", model="", optional_params={}, litellm_params={}) + assert url == expected + + +def test_validate_environment_uses_bearer_auth(config, api_key): + headers = config.validate_environment(headers={}, model="", messages=[], optional_params={}, litellm_params={}) + assert headers == {"Authorization": f"Bearer {api_key}"} + + +def test_upload_request_is_multipart_with_batch_purpose(config): + body = config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"), + optional_params={}, + litellm_params={}, + ) + assert body == { + "file": ("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), + "purpose": (None, "batch"), + } + + +@pytest.mark.parametrize( + "openai_purpose,mistral_purpose", + [("batch", "batch"), ("fine-tune", "fine-tune"), ("ocr", "ocr"), ("assistants", "batch"), ("user_data", "batch")], +) +def test_upload_request_maps_purpose_onto_mistral_enum(config, openai_purpose, mistral_purpose): + body = config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=openai_purpose), + optional_params={}, + litellm_params={}, + ) + assert body["purpose"] == (None, mistral_purpose) + + +def test_upload_request_requires_file(config): + with pytest.raises(ValueError, match="File data is required"): + config.transform_create_file_request( + model="", create_file_data=CreateFileRequest(purpose="batch"), optional_params={}, litellm_params={} + ) + + +def test_upload_response_maps_onto_openai_file_object(config): + obj = config.transform_create_file_response( + model=None, raw_response=_response(_file()), logging_obj=None, litellm_params={} + ) + assert obj == OpenAIFileObject( + id=FILE_ID, + bytes=13000, + created_at=1_716_963_433, + filename="batch_input.jsonl", + object="file", + purpose="batch", + status="uploaded", + ) + + +def test_file_response_with_ocr_purpose_maps_onto_user_data(config): + obj = config.transform_retrieve_file_response( + raw_response=_response(_file(purpose="ocr", expires_at=1_800_000_000)), logging_obj=None, litellm_params={} + ) + assert obj.purpose == "user_data" + assert obj.expires_at == 1_800_000_000 + + +@pytest.mark.parametrize( + "method,suffix", + [ + ("transform_retrieve_file_request", ""), + ("transform_delete_file_request", ""), + ], +) +def test_single_file_urls_encode_id_and_honor_api_base(config, method, suffix): + url, params = getattr(config, method)( + file_id="id/with slash", optional_params={}, litellm_params={"api_base": "https://mistral.internal/v1"} + ) + assert url == f"https://mistral.internal/v1/files/id%2Fwith%20slash{suffix}" + assert params == {} + + +def test_file_content_url(config): + url, params = config.transform_file_content_request( + file_content_request=FileContentRequest(file_id=FILE_ID), optional_params={}, litellm_params={} + ) + assert url == f"https://api.mistral.ai/v1/files/{FILE_ID}/content" + assert params == {} + + +def test_file_content_response_is_binary_passthrough(config): + raw = httpx.Response( + 200, content=b'{"custom_id":"0","response":{"status_code":200}}\n', request=httpx.Request("GET", "https://x") + ) + out = config.transform_file_content_response(raw_response=raw, logging_obj=None, litellm_params={}) + assert out.content == b'{"custom_id":"0","response":{"status_code":200}}\n' + + +def test_delete_response(config): + out = config.transform_delete_file_response( + raw_response=_response({"id": FILE_ID, "object": "file", "deleted": True}), logging_obj=None, litellm_params={} + ) + assert out == FileDeleted(id=FILE_ID, deleted=True, object="file") + + +def test_list_request_filters_by_mapped_purpose(config): + url, params = config.transform_list_files_request(purpose="batch", optional_params={}, litellm_params={}) + assert url == "https://api.mistral.ai/v1/files" + assert params == {"purpose": "batch"} + _, no_params = config.transform_list_files_request(purpose=None, optional_params={}, litellm_params={}) + assert no_params == {} + + +def test_list_response(config): + out = config.transform_list_files_response( + raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}), + logging_obj=None, + litellm_params={}, + ) + assert [f.id for f in out] == [FILE_ID, "second"] + assert out[1].filename == "b.jsonl" diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 40e54f71eeb..9fe6f38003f 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -72,9 +72,11 @@ def test_ocr3_pricing_entry(cost_map_path: Path) -> None: assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}" assert info["litellm_provider"] == "mistral" assert info["mode"] == "ocr" - assert info["supported_endpoints"] == ["/v1/ocr"] + assert info["supported_endpoints"] == ["/v1/ocr", "/v1/batch"] assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE + assert info["ocr_cost_per_page_batches"] == OCR3_COST_PER_PAGE / 2 + assert info["annotation_cost_per_page_batches"] == OCR3_ANNOTATION_COST_PER_PAGE / 2 def test_ocr3_model_info_price(local_model_cost_map) -> None: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index e42608c9904..61739846706 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -992,7 +992,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, "input_dbu_cost_per_token": {"type": "number"}, "annotation_cost_per_page": {"type": "number"}, + "annotation_cost_per_page_batches": {"type": "number"}, "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_page_batches": {"type": "number"}, "ocr_cost_per_credit": {"type": "number"}, "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, From 2e5f5a95c813b940dc7f654c10b2ce2036c6fb2a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:14:16 -0400 Subject: [PATCH 016/349] fix(proxy): retrieve model-routed file ids from the deployment's provider GET /v1/files/{id} for an id encoded with a non-OpenAI deployment forwarded the deployment credentials but let custom_llm_provider default to openai, so a Mistral file was fetched from api.openai.com with the Mistral key and 401'd. Delete and content already passed the provider through; retrieve now does too. --- .../openai_files_endpoints/files_endpoints.py | 5 +- .../test_files_endpoint.py | 59 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c315d30b8f3..49a01495d65 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1139,7 +1139,10 @@ async def get_file( include_internal_credentials=True, ) - response = await litellm.afile_retrieve(**data) + response = await litellm.afile_retrieve( + custom_llm_provider=credentials["custom_llm_provider"], + **data, + ) # Keep the encoded ID in response if it was originally encoded if original_file_id and response and hasattr(response, "id") and response.id: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5f1e7e1fe0c..5faae166fca 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4819,3 +4819,62 @@ def test_get_file_content_keeps_the_status_of_a_rejection_raised_inside_the_rout error = response.json()["error"] assert error["message"].startswith("Storage backend error") assert (error["type"], error["param"], error["code"]) == ("invalid_request_error", "file_id", "400") + + +def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFixture, monkeypatch): + """ + Regression: a file id encoded with a non-OpenAI deployment (here Mistral) must be + retrieved from that deployment's provider. Before the fix the retrieve path only + forwarded the credentials and let ``custom_llm_provider`` default to openai, so a + Mistral file id was sent to api.openai.com with the Mistral key and 401'd. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"}, + "model_info": {"id": "mistral-ocr-id"}, + } + ] + ) + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_retrieve(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", + object="file", + bytes=2, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["custom_llm_provider"] == "mistral" + assert captured_kwargs["api_key"] == "mistral-key" + assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df" + assert response.json()["id"] == encoded_id From c246f75e3ec93192f95e0cb2fc3f50485bf13ebc Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:37:01 -0400 Subject: [PATCH 017/349] refactor(mistral): satisfy type-discipline and basedpyright gates for files/batches configs --- litellm/cost_calculator.py | 23 ++-- litellm/files/main.py | 4 +- .../llms/mistral/batches/transformation.py | 114 ++++++++++------ litellm/llms/mistral/common_utils.py | 13 +- litellm/llms/mistral/files/transformation.py | 129 +++++++++++------- .../openai_files_endpoints/files_endpoints.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 46 +++++-- tests/test_litellm/batches/test_main.py | 76 +++-------- .../test_mistral_batches_transformation.py | 16 ++- .../test_mistral_files_transformation.py | 8 +- 10 files changed, 251 insertions(+), 180 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0c0c6b05df8..7c95941d77d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2006,13 +2006,11 @@ def ocr_batch_cost( ``ocr_cost``. """ has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) - if has_ocr_pricing: - resolved_info: ModelInfo | None = model_info - else: - try: - resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - resolved_info = None + resolved_info: Final = ( + model_info + if has_ocr_pricing + else _lookup_model_info_or_none(model=model, custom_llm_provider=custom_llm_provider) + ) if resolved_info is None: verbose_logger.warning( "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", @@ -2022,9 +2020,7 @@ def ocr_batch_cost( return 0.0, 0.0 page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") - annotation_rate: Final = _first_price( - resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page" - ) + annotation_rate: Final = _first_price(resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page") pages_processed: Final = usage_info.pages_processed or 0 annotation_pages: Final = usage_info.pages_processed_annotation or 0 if page_rate is None and pages_processed > 0: @@ -2039,6 +2035,13 @@ def ocr_batch_cost( return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0 +def _lookup_model_info_or_none(model: str, custom_llm_provider: str | None) -> ModelInfo | None: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + return None + + def _first_price(model_info: ModelInfo, *keys: str) -> float | None: return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) diff --git a/litellm/files/main.py b/litellm/files/main.py index 3d90bf4f299..5cf0f8e576a 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -32,9 +32,7 @@ FileCreateProvider = Literal[ FileRetrieveProvider = Literal[ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" ] -FileDeleteProvider = Literal[ - "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral" -] +FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"] FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] import litellm from litellm import get_secret_str diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py index 399319e590b..ef9ee5ff503 100644 --- a/litellm/llms/mistral/batches/transformation.py +++ b/litellm/llms/mistral/batches/transformation.py @@ -7,14 +7,16 @@ Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_ so the shared batch cost accounting reads them without a provider branch. """ +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, Literal +from typing import Final, Literal, TypeAlias import httpx from openai.types.batch import BatchRequestCounts from openai.types.batch import Errors as BatchErrors from openai.types.batch_error import BatchError from pydantic import BaseModel, ConfigDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig @@ -24,13 +26,14 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error -MistralBatchStatus = Literal[ +MistralBatchStatus: TypeAlias = Literal[ "QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED" ] -OpenAIBatchStatus = Literal[ +OpenAIBatchStatus: TypeAlias = Literal[ "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" ] +_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) # mutable-ok: frozen at module scope _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType( { "QUEUED": "validating", @@ -44,6 +47,23 @@ _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = Ma ) +class MistralCreateBatchJobRequest(TypedDict): + """Body of ``POST /v1/batch/jobs``.""" + + input_files: ReadOnly[tuple[str, ...]] + endpoint: ReadOnly[str] + model: ReadOnly[str] + metadata: NotRequired[ReadOnly[Mapping[str, str]]] + + +class MistralPresignedRequest(TypedDict): + """A fully-formed request the shared HTTP handler sends as-is (its ``method`` branch).""" + + method: ReadOnly[Literal["GET"]] + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + + class MistralBatchError(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -69,7 +89,18 @@ class MistralBatchJob(BaseModel): output_file: str | None = None error_file: str | None = None errors: tuple[MistralBatchError, ...] = () - metadata: dict[str, str] | None = None + metadata: dict[str, str] | None = None # mutable-ok: LiteLLMBatch.metadata is typed as dict + + +def _to_batch_errors(errors: Sequence[MistralBatchError]) -> BatchErrors | None: + if not errors: + return None + return BatchErrors( + object="list", + data=[ # mutable-ok: openai Batch.Errors.data is typed as list + BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in errors + ], + ) def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: @@ -90,14 +121,7 @@ def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: cancelled_at=terminal_at if status == "cancelled" else None, output_file_id=job.output_file, error_file_id=job.error_file, - errors=( - BatchErrors( - object="list", - data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors], - ) - if job.errors - else None - ), + errors=_to_batch_errors(job.errors), request_counts=BatchRequestCounts( total=job.total_requests, completed=job.succeeded_requests, @@ -114,14 +138,14 @@ class MistralBatchesConfig(BaseBatchesConfig): def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: # mutable-ok: BaseBatchesConfig signature return get_mistral_auth_headers(headers, api_key) def get_complete_batch_url( @@ -129,8 +153,8 @@ class MistralBatchesConfig(BaseBatchesConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], data: CreateBatchRequest, ) -> str: return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" @@ -139,48 +163,58 @@ class MistralBatchesConfig(BaseBatchesConfig): self, model: str, create_batch_data: CreateBatchRequest, - optional_params: dict, - litellm_params: dict, - ) -> dict[str, object]: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature + input_file_id: Final = create_batch_data.get("input_file_id") + endpoint: Final = create_batch_data.get("endpoint") + if input_file_id is None or endpoint is None: + raise ValueError("input_file_id and endpoint are required to create a Mistral batch job") metadata: Final = create_batch_data.get("metadata") - return { - "input_files": [create_batch_data["input_file_id"]], - "endpoint": create_batch_data["endpoint"], - "model": model, - **({"metadata": metadata} if metadata else {}), - **(create_batch_data.get("extra_body") or {}), - } + body: Final = ( + MistralCreateBatchJobRequest( + input_files=(input_file_id,), endpoint=endpoint, model=model, metadata=metadata + ) + if metadata + else MistralCreateBatchJobRequest(input_files=(input_file_id,), endpoint=endpoint, model=model) + ) + return dict(body) # mutable-ok: BaseBatchesConfig signature def transform_create_batch_response( self, model: str | None, raw_response: httpx.Response, logging_obj: object, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> LiteLLMBatch: return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) def transform_retrieve_batch_request( self, batch_id: str, - optional_params: dict, - litellm_params: dict, - ) -> dict[str, object]: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - return { - "method": "GET", - "url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}", - "headers": get_mistral_auth_headers({}, litellm_params.get("api_key")), - } + api_base: Final = litellm_params.get("api_base") + api_key: Final = litellm_params.get("api_key") + request: Final = MistralPresignedRequest( + method="GET", + url=f"{get_mistral_api_base(api_base if isinstance(api_base, str) else None)}/v1/batch/jobs/{encoded_batch_id}", + headers=get_mistral_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None), + ) + return dict(request) # mutable-ok: BaseBatchesConfig signature def transform_retrieve_batch_response( self, model: str | None, raw_response: httpx.Response, logging_obj: object, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> LiteLLMBatch: return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: return mistral_error(error_message, status_code, headers) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py index 9ea501c860d..2f14328afdf 100644 --- a/litellm/llms/mistral/common_utils.py +++ b/litellm/llms/mistral/common_utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -19,18 +20,22 @@ def get_mistral_api_base(api_base: str | None) -> str: return resolved.removesuffix("/v1") -def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict: +def get_mistral_auth_headers( + headers: Mapping[str, str], api_key: str | None +) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR) if resolved_key is None: raise ValueError( "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" ) - return {**headers, "Authorization": f"Bearer {resolved_key}"} + return dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict -def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError: +def mistral_error(error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers) -> MistralError: return MistralError( status_code=status_code, message=error_message, - headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), + headers=headers + if isinstance(headers, httpx.Headers) + else httpx.Headers(dict(headers)), # mutable-ok: httpx.Headers takes a dict ) diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 071b6f58569..6d58311813c 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -7,11 +7,13 @@ Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes. """ import time -from typing import Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, TypeAlias import httpx from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -29,7 +31,16 @@ from litellm.types.utils import LlmProviders from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error -MistralFilePurpose = Literal["fine-tune", "batch", "ocr"] +MistralFilePurpose: TypeAlias = Literal["fine-tune", "batch", "ocr"] + +_NO_QUERY_PARAMS: Final[dict[str, str]] = {} # mutable-ok: BaseFilesConfig request transforms return tuple[str, dict] + + +class MistralMultipartUpload(TypedDict): + """``files=`` payload for ``POST /v1/files``: each value is an httpx multipart tuple.""" + + file: ReadOnly[tuple[str, object, str]] + purpose: ReadOnly[tuple[None, MistralFilePurpose]] class MistralFile(BaseModel): @@ -85,6 +96,11 @@ def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: return "batch" +def _api_base_from(litellm_params: Mapping[str, object]) -> str: + api_base: Final = litellm_params.get("api_base") + return get_mistral_api_base(api_base if isinstance(api_base, str) else None) + + class MistralFilesConfig(BaseFilesConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -95,99 +111,103 @@ class MistralFilesConfig(BaseFilesConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: return f"{get_mistral_api_base(api_base)}/v1/files" - def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str: + def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}" + return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: return mistral_error(error_message, status_code, headers) def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list, - optional_params: dict, - litellm_params: dict, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: # mutable-ok: BaseFilesConfig signature return get_mistral_auth_headers(headers, api_key) - def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: - return ["purpose"] + def get_supported_openai_params( + self, model: str + ) -> list[OpenAICreateFileRequestOptionalParams]: # mutable-ok: BaseFilesConfig signature + return ["purpose"] # mutable-ok: BaseFilesConfig signature def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: Mapping[str, object], + optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: # mutable-ok: BaseConfig signature return optional_params def transform_create_file_request( self, model: str, create_file_data: CreateFileRequest, - optional_params: dict, - litellm_params: dict, - ) -> dict: - file_data: Final = create_file_data.get("file") - if file_data is None: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseFilesConfig signature + if "file" not in create_file_data: raise ValueError("File data is required") - extracted: Final = extract_file_data(file_data) + extracted: Final = extract_file_data(create_file_data["file"]) filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" content_type: Final = extracted.get("content_type") or "application/octet-stream" - return { - "file": (filename, extracted["content"], content_type), - "purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), - } + upload: Final = MistralMultipartUpload( + file=(filename, extracted["content"], content_type), + purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + ) + return dict(upload) # mutable-ok: BaseFilesConfig signature def transform_create_file_response( self, model: str | None, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) def transform_retrieve_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_id, litellm_params), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS def transform_retrieve_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) def transform_delete_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_id, litellm_params), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS def transform_delete_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> FileDeleted: deleted: Final = MistralFileDeleted.model_validate(raw_response.json()) return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") @@ -195,32 +215,39 @@ class MistralFilesConfig(BaseFilesConfig): def transform_list_files_request( self, purpose: str | None, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {} - return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + url: Final = f"{_api_base_from(litellm_params)}/v1/files" + if not purpose: + return url, _NO_QUERY_PARAMS + return url, {"purpose": _to_mistral_purpose(purpose)} # mutable-ok: BaseFilesConfig signature returns dict def transform_list_files_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - ) -> list[OpenAIFileObject]: - return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data] + litellm_params: Mapping[str, object], + ) -> list[OpenAIFileObject]: # mutable-ok: BaseFilesConfig signature + return [ # mutable-ok: BaseFilesConfig signature + _to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data + ] def transform_file_content_request( self, file_content_request: FileContentRequest, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + file_id: Final = file_content_request.get("file_id") + if file_id is None: + raise ValueError("file_id is required to download file content") + return self._file_url(file_id, litellm_params, suffix="/content"), _NO_QUERY_PARAMS def transform_file_content_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> HttpxBinaryResponseContent: return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 49a01495d65..b3bb1fa9a01 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1130,7 +1130,7 @@ async def get_file( check_file_id_encoding=True, ) - if should_route: + if should_route and credentials is not None: # Use model-based routing with credentials from config prepare_data_with_credentials( data=data, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 85f267c2db0..56cd3298db6 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -645,9 +645,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): lambda content, model: pytest.fail("raw vertex path should not run"), ) - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[], custom_llm_provider="vertex_ai" - ) + result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") assert result.cost == 0.0 assert result.usage.total_tokens == 0 assert result.models == [] @@ -1250,6 +1248,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): result set - zero cost, zero usage, no models - instead of letting the file fetch raise "Output file id is None" on every aretrieve_batch logging poll. """ + # The output-file fetch must not even be attempted when there is no output file. async def _must_not_fetch(*args, **kwargs): pytest.fail("_fetch_batch_output_file_content should not be called") @@ -1376,7 +1375,10 @@ def test_anthropic_response_body_is_result_message(): def test_anthropic_usage_conversion_includes_cache_tokens(): - body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)} + body = { + "model": "claude-sonnet-4-5-20250929", + "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000), + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic") assert usage.prompt_tokens == 11000 assert usage.completion_tokens == 200 @@ -1391,7 +1393,9 @@ def test_bedrock_model_output_line_success_check(): "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True - assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + assert ( + bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + ) def test_bedrock_cost_uses_deployment_model_name(): @@ -1445,7 +1449,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): rows = [ { "custom_id": "req-1", - "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, + "response": { + "status_code": 200, + "body": { + "model": "gpt-5.2", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, } ] result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") @@ -1487,7 +1497,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"), ) - result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic") + result = bu._aggregate_batch_cost_usage_models( + entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic" + ) assert result.cost == pytest.approx(0.3) assert seen[0]["model"] == "claude-sonnet-4-5-20250929" @@ -1522,7 +1534,11 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): ) assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 11000, + 200, + 11200, + ) assert result.models == ["claude-sonnet-4-5"] @@ -1689,7 +1705,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> def test_bedrock_converse_shaped_batch_usage_is_parsed(): - body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}} + body = { + "model": "us.amazon.nova-lite-v1:0", + "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}, + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742) @@ -1739,6 +1758,7 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): # batch_cost_is_final # --------------------------------------------------------------------------- # + def _retrieved_batch( status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None ) -> LiteLLMBatch: @@ -1798,7 +1818,9 @@ def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest") usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: usage_info["pages_processed_annotation"] = annotation_pages - return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info) + return _success_row( + model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info + ) def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch): @@ -1835,7 +1857,9 @@ def test_ocr_rows_bill_annotation_pages_separately(monkeypatch): "annotation_cost_per_page_batches": 0.0025, }, ) - result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral") + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral" + ) assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025) diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c5a33dd6508..26dc4083b0b 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -66,9 +66,7 @@ def seams(): stack.enter_context(patch.object(bm, "openai_batches_instance", openai_i)) stack.enter_context(patch.object(bm, "azure_batches_instance", azure_i)) stack.enter_context(patch.object(bm, "vertex_ai_batches_instance", vertex_i)) - stack.enter_context( - patch.object(bm, "anthropic_batches_instance", anthropic_i) - ) + stack.enter_context(patch.object(bm, "anthropic_batches_instance", anthropic_i)) stack.enter_context(patch.object(bm, "base_llm_http_handler", base_http)) stack.enter_context(patch.object(bm, "BedrockBatchesHandler", bedrock_arn)) yield Seams( @@ -174,9 +172,7 @@ def test_create__provider_config_routes_to_base_http_handler(seams): "get_provider_batches_config", return_value=MagicMock(name="provider_config"), ): - result = bm.create_batch( - **CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model" - ) + result = bm.create_batch(**CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model") assert result is seams.base_http.create_batch.return_value _assert_only(seams.base_http.create_batch, seams, "create_batch") @@ -281,9 +277,7 @@ def test_retrieve__bedrock_model_invocation_job_arn(seams): result = bm.retrieve_batch(batch_id=arn, custom_llm_provider="bedrock") seams.bedrock_arn._handle_model_invocation_job_status.assert_called_once() - assert ( - result is seams.bedrock_arn._handle_model_invocation_job_status.return_value - ) + assert result is seams.bedrock_arn._handle_model_invocation_job_status.return_value seams.bedrock_arn._handle_async_invoke_status.assert_not_called() @@ -385,9 +379,7 @@ def test_cancel__unsupported_provider_raises_badrequest(seams): def test_cancel__async_flag_propagates_is_async(seams): - bm.cancel_batch( - batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True - ) + bm.cancel_batch(batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True) assert seams.openai.cancel_batch.call_args.kwargs["_is_async"] is True @@ -415,9 +407,7 @@ async def test_acreate_batch_delegates_to_create_batch(): @pytest.mark.asyncio async def test_aretrieve_batch_delegates_to_retrieve_batch(): with patch.object(bm, "retrieve_batch", MagicMock(return_value="SENTINEL")) as m: - result = await bm.aretrieve_batch( - batch_id="batch-1", custom_llm_provider="azure" - ) + result = await bm.aretrieve_batch(batch_id="batch-1", custom_llm_provider="azure") assert result == "SENTINEL" assert m.call_count == 1 @@ -429,9 +419,7 @@ async def test_aretrieve_batch_delegates_to_retrieve_batch(): @pytest.mark.asyncio async def test_alist_batches_delegates_to_list_batches(): with patch.object(bm, "list_batches", MagicMock(return_value="SENTINEL")) as m: - result = await bm.alist_batches( - after="cur", limit=3, custom_llm_provider="vertex_ai" - ) + result = await bm.alist_batches(after="cur", limit=3, custom_llm_provider="vertex_ai") assert result == "SENTINEL" assert m.call_count == 1 @@ -444,9 +432,7 @@ async def test_alist_batches_delegates_to_list_batches(): @pytest.mark.asyncio async def test_acancel_batch_delegates_to_cancel_batch(): with patch.object(bm, "cancel_batch", MagicMock(return_value="SENTINEL")) as m: - result = await bm.acancel_batch( - batch_id="batch-1", custom_llm_provider="openai" - ) + result = await bm.acancel_batch(batch_id="batch-1", custom_llm_provider="openai") assert result == "SENTINEL" assert m.call_count == 1 @@ -499,9 +485,7 @@ def _sent(mock_method, *keys): def test_create__openai_credentials_passthrough(seams): bm.create_batch(**CREATE_KW, custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries" - ) == { + assert _sent(seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -512,9 +496,7 @@ def test_create__openai_credentials_passthrough(seams): def test_create__azure_credentials_passthrough(seams): bm.create_batch(**CREATE_KW, custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.create_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.create_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -564,9 +546,7 @@ def test_create__provider_config_credentials_passthrough(seams): def test_retrieve__openai_credentials_passthrough(seams): bm.retrieve_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.retrieve_batch, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.retrieve_batch, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -576,9 +556,7 @@ def test_retrieve__openai_credentials_passthrough(seams): def test_retrieve__azure_credentials_passthrough(seams): bm.retrieve_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.retrieve_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.retrieve_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -640,9 +618,7 @@ def test_retrieve__provider_config_credentials_passthrough(seams): def test_list__openai_credentials_passthrough(seams): bm.list_batches(custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.list_batches, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.list_batches, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -652,9 +628,7 @@ def test_list__openai_credentials_passthrough(seams): def test_list__azure_credentials_passthrough(seams): bm.list_batches(custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.list_batches, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.list_batches, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -682,9 +656,7 @@ def test_list__vertex_credentials_passthrough(seams): def test_cancel__openai_credentials_passthrough(seams): bm.cancel_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.cancel_batch, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.cancel_batch, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -694,9 +666,7 @@ def test_cancel__openai_credentials_passthrough(seams): def test_cancel__azure_credentials_passthrough(seams): bm.cancel_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.cancel_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.cancel_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -786,18 +756,16 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams): - with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg: - result = bm.create_batch( - completion_window="24h", - endpoint="/v1/ocr", - input_file_id="file-abc", - custom_llm_provider="mistral", - model="mistral/mistral-ocr-latest", - ) + result = bm.create_batch( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-abc", + custom_llm_provider="mistral", + model="mistral/mistral-ocr-latest", + ) assert result is seams.base_http.create_batch.return_value _assert_only(seams.base_http.create_batch, seams, "create_batch") - get_cfg.assert_called_once() forwarded = seams.base_http.create_batch.call_args.kwargs assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" assert forwarded["model"] == "mistral-ocr-latest" diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py index 03e9c351a30..03cfeedece2 100644 --- a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -92,26 +92,34 @@ def test_create_request_maps_openai_fields_onto_mistral_job(config): model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={} ) assert body == { - "input_files": ["file-123"], + "input_files": ("file-123",), "endpoint": "/v1/ocr", "model": "mistral-ocr-latest", "metadata": {"team": "docs"}, } -def test_create_request_omits_empty_metadata_and_forwards_extra_body(config): +def test_create_request_omits_empty_metadata(config): data = CreateBatchRequest( completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-123", metadata=None, - extra_body={"timeout_hours": 48}, ) body = config.transform_create_batch_request( model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={} ) assert "metadata" not in body - assert body["timeout_hours"] == 48 + + +def test_create_request_requires_input_file_and_endpoint(config): + with pytest.raises(ValueError, match="input_file_id and endpoint are required"): + config.transform_create_batch_request( + model="m", + create_batch_data=CreateBatchRequest(completion_window="24h"), + optional_params={}, + litellm_params={}, + ) @pytest.mark.parametrize( diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index d6ad8a34b35..f62645be7ee 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -79,7 +79,9 @@ def test_validate_environment_uses_bearer_auth(config, api_key): def test_upload_request_is_multipart_with_batch_purpose(config): body = config.transform_create_file_request( model="", - create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"), + create_file_data=CreateFileRequest( + file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch" + ), optional_params={}, litellm_params={}, ) @@ -181,7 +183,9 @@ def test_list_request_filters_by_mapped_purpose(config): def test_list_response(config): out = config.transform_list_files_response( - raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}), + raw_response=_response( + {"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2} + ), logging_obj=None, litellm_params={}, ) From 91e7df3d8fe83c37268b61080a37e48ef0e63e3f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 13:35:35 -0400 Subject: [PATCH 018/349] chore: regenerate cost-map schema and UI API types, drop test banner comments --- model_prices_and_context_window.schema.json | 8 ++++++++ tests/test_litellm/batches/test_batch_utils.py | 5 ----- .../batches/test_mistral_batches_transformation.py | 10 ---------- .../llms/mistral/ocr/test_mistral_ocr_cost.py | 1 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++++++++ 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 47a1934a703..09385db728b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -53,6 +53,10 @@ "type": "number", "minimum": 0 }, + "annotation_cost_per_page_batches": { + "type": "number", + "minimum": 0 + }, "audio_transcription_config": { "type": "string" }, @@ -432,6 +436,10 @@ "type": "number", "minimum": 0 }, + "ocr_cost_per_page_batches": { + "type": "number", + "minimum": 0 + }, "output_cost_per_audio_token": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 56cd3298db6..0d66c0eb5ec 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1809,11 +1809,6 @@ class TestBatchCostIsFinal: assert bu.batch_cost_is_final(_retrieved_batch(status)) is True -# =========================================================================== # -# OCR batch output lines (Mistral /v1/ocr batches) are billed per page, not per token -# =========================================================================== # - - def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py index 03cfeedece2..4073879e3b8 100644 --- a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -76,11 +76,6 @@ def test_custom_llm_provider(config): assert config.custom_llm_provider == LlmProviders.MISTRAL -# --------------------------------------------------------------------------- # -# create -# --------------------------------------------------------------------------- # - - def test_create_request_maps_openai_fields_onto_mistral_job(config): data = CreateBatchRequest( completion_window="24h", @@ -175,11 +170,6 @@ def test_create_response_maps_job_onto_openai_batch(config): assert batch.metadata == {"job_type": "testing"} -# --------------------------------------------------------------------------- # -# retrieve -# --------------------------------------------------------------------------- # - - def test_retrieve_request_is_presigned_get_with_auth(config, api_key): req = config.transform_retrieve_batch_request( batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"} diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 9fe6f38003f..d72e866949f 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -63,7 +63,6 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) - @pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) def test_ocr3_pricing_entry(cost_map_path: Path) -> None: with open(cost_map_path) as f: diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0f0c1fc9af4..5d1ed79098e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29493,6 +29493,8 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Annotation Cost Per Page Batches */ + annotation_cost_per_page_batches?: number | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -29704,6 +29706,8 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Ocr Cost Per Page Batches */ + ocr_cost_per_page_batches?: number | null; /** Organization */ organization?: string | null; /** Otpm */ @@ -39684,6 +39688,8 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Annotation Cost Per Page Batches */ + annotation_cost_per_page_batches?: number | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -39895,6 +39901,8 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Ocr Cost Per Page Batches */ + ocr_cost_per_page_batches?: number | null; /** Organization */ organization?: string | null; /** Otpm */ From edd5727f3c9077f449eec37367ace43945e649fa Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 17:28:27 -0400 Subject: [PATCH 019/349] fix(proxy): enforce key/team/org/project model grants on model-routed file and batch credentials Files and batches routes take their model from a header, query param or a model-encoded resource id, which the auth layer never sees, so any key could name any deployment and act on that provider account with its server-side key. Every caller-supplied model now goes through can_key_call_resolved_model before deployment credentials are resolved, covering file create/retrieve/content/ delete/list, batch create/retrieve/list/cancel, and vector store files. --- litellm/proxy/batches_endpoints/endpoints.py | 17 +- .../openai_files_endpoints/common_utils.py | 61 +++++- .../openai_files_endpoints/files_endpoints.py | 20 +- .../vector_store_files_endpoints/endpoints.py | 6 +- .../proxy/batches_endpoints/test_endpoints.py | 58 +++++- .../test_files_endpoint.py | 184 ++++++++++++++++-- .../test_batch_x_litellm_model_encoding.py | 53 ++--- 7 files changed, 322 insertions(+), 77 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..c99f66d032e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -34,9 +34,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, + get_authorized_credentials_for_model, get_batch_from_database, get_batch_id_from_unified_batch_id, - get_credentials_for_model, get_model_id_from_unified_batch_id, get_models_from_unified_file_id, get_original_file_id, @@ -218,9 +218,10 @@ async def create_batch( # SCENARIO 1: File ID is encoded with model info if model_from_file_id is not None and input_file_id: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_file_id, + user_api_key_dict=user_api_key_dict, operation_context="batch creation (file created with model)", ) @@ -310,9 +311,10 @@ async def create_batch( # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback if model_param: # SCENARIO 2: Use model-based routing from header/query/body - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, + user_api_key_dict=user_api_key_dict, operation_context="batch creation", ) @@ -540,9 +542,10 @@ async def retrieve_batch( # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context="batch retrieval (batch created with model)", ) @@ -764,9 +767,10 @@ async def list_batches( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ): # SCENARIO 2: Use model-based routing from header/query/body - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, + user_api_key_dict=user_api_key_dict, operation_context="batch listing", ) @@ -952,9 +956,10 @@ async def cancel_batch( # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context="batch cancellation (batch created with model)", ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b1f282a0978..4202a6d1689 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -351,6 +351,10 @@ def get_credentials_for_model( """ Retrieve API credentials for a model from the LLM Router. + Does not check whether the caller may use ``model_id``; use + ``get_authorized_credentials_for_model`` for anything driven by a caller-supplied + model name (request body, header, query param, or a model-encoded resource id). + Args: llm_router: LiteLLM Router instance model_id: Model name or deployment ID @@ -381,6 +385,48 @@ def get_credentials_for_model( return credentials +async def authorize_model_for_key( + model_id: str, + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", +) -> None: + """ + Enforce the caller's model grants on a model name the auth layer never saw. + + The files and batches routes carry their model in a header, query param, or a + model-encoded resource id rather than the request body, so ``user_api_key_auth`` + cannot check it. Run the same key, team (incl. team-member and access-group + fallbacks), org and project allowlist checks a chat request would get, so a + restricted key cannot borrow another deployment's server-side credentials. + + Raises: + ProxyException (403): the caller is not allowed to use ``model_id`` + """ + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + await can_key_call_resolved_model( + model=model_id, + llm_model_list=None, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + + +async def get_authorized_credentials_for_model( + llm_router: Optional["Router"], + model_id: str, + user_api_key_dict: "UserAPIKeyAuth", + operation_context: str = "file operation", +) -> dict: # mutable-ok: same contract as get_credentials_for_model, callers merge it into request data + """``get_credentials_for_model`` gated by ``authorize_model_for_key``.""" + await authorize_model_for_key(model_id=model_id, llm_router=llm_router, user_api_key_dict=user_api_key_dict) + return get_credentials_for_model( + llm_router=llm_router, + model_id=model_id, + operation_context=operation_context, + ) + + def get_team_provider_credentials( llm_router: Optional["Router"], user_api_key_dict: "UserAPIKeyAuth", @@ -573,21 +619,27 @@ def prepare_data_with_credentials( data["file_id"] = file_id -def handle_model_based_routing( +async def handle_model_based_routing( file_id: str, request, # FastAPI Request object llm_router, # Router instance data: dict, + user_api_key_dict: "UserAPIKeyAuth", check_file_id_encoding: bool = True, ) -> tuple[bool, str | None, str | None, dict | None]: """ Orchestrate model-based credential routing for file operations. + The model name comes from the caller (embedded in the file id, or a header, query + param or body field), so it is authorized against the caller's key, team, org and + project grants before any deployment credentials are resolved. + Args: file_id: File ID (may contain embedded model info) request: FastAPI request object llm_router: LiteLLM Router instance data: Request data dictionary + user_api_key_dict: The authenticated caller check_file_id_encoding: Whether to check for embedded model in file_id Returns: @@ -599,6 +651,7 @@ def handle_model_based_routing( Raises: HTTPException: If router unavailable or model not found + ProxyException: If the caller is not allowed to use the model """ model_from_id, model_from_param = extract_model_from_sources( file_id=file_id, @@ -608,9 +661,10 @@ def handle_model_based_routing( # Priority 1: Model embedded in file_id if check_file_id_encoding and model_from_id is not None: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context=f"file operation (file created with model '{model_from_id}')", ) original_file_id: Final = get_original_file_id(file_id) @@ -618,9 +672,10 @@ def handle_model_based_routing( # Priority 2: Model from header/query/body elif model_from_param is not None: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_param, + user_api_key_dict=user_api_key_dict, operation_context="file operation", ) return True, model_from_param, None, credentials diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b3bb1fa9a01..0efd618e171 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -68,7 +68,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( apply_team_provider_credentials, encode_file_id_with_model, extract_file_creation_params, - get_credentials_for_model, + get_authorized_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, validate_file_list_limit, @@ -267,9 +267,10 @@ async def route_create_file( # NEW: Handle model-based routing (no DB required) if model is not None: # Get credentials from model_list via router - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model, + user_api_key_dict=user_api_key_dict, operation_context="file upload", ) @@ -907,11 +908,12 @@ async def get_file_content( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1122,11 +1124,12 @@ async def get_file( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1330,11 +1333,12 @@ async def delete_file( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1517,11 +1521,12 @@ async def list_files( response: Any | None = None # Check for model-based credential routing (no file_id encoding check for list) - should_route, model_used, _, credentials = handle_model_based_routing( + should_route, model_used, _, credentials = await handle_model_based_routing( file_id="", # No file_id for list endpoint request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=False, ) @@ -1548,9 +1553,10 @@ async def list_files( status_code=500, detail="LLM Router not initialized. Ensure models added to proxy.", ) - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=target_model_names_list[0], + user_api_key_dict=user_api_key_dict, operation_context="file list", ) prepare_data_with_credentials(data=data, credentials=credentials) diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 957ed9fd0b9..11ef8efb598 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -144,11 +144,12 @@ async def _update_request_data_with_managed_file_id( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -273,11 +274,12 @@ async def _update_request_data_with_model_routing_hint( _model_used, _original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id="", request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=False, ) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..57e42e79a42 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -177,6 +177,8 @@ def harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1161,6 +1163,8 @@ def retrieve_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1616,6 +1620,8 @@ def list_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2012,6 +2018,8 @@ def cancel_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2733,8 +2741,6 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc assert cancel_harness.router_acancel.call_count == 1 - - @pytest.mark.asyncio async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness): with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): @@ -2762,3 +2768,51 @@ async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retriev metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {} assert metadata.get("batch_ignore_default_logging") is None + + +def _key_restricted_to(*models: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-restricted", team_id="team-a", team_models=list(models), models=list(models)) + + +@pytest.mark.asyncio +async def test_create__header_model_rejects_key_without_model_grant(harness): + """A key not granted the model named in x-litellm-model must not receive that deployment's credentials.""" + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness, user=_key_restricted_to("azure/gpt-4o"), headers={"x-litellm-model": "vertex-model"}) + + assert exc_info.value.code == "403" + harness.creds_resolver.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__header_model_allows_key_with_model_grant(harness): + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + await call_create(harness, user=_key_restricted_to("vertex-model"), headers={"x-litellm-model": "vertex-model"}) + + harness.creds_resolver.assert_called_once_with(model_id="vertex-model") + assert harness.acreate_kwargs()["custom_llm_provider"] == "vertex_ai" + + +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id_rejects_key_without_model_grant(retrieve_harness): + """The model embedded in a batch id is caller-controlled, so it is checked against the key's grants too.""" + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + retrieve_harness.creds_resolver.assert_not_called() + retrieve_harness.litellm_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_harness): + with pytest.raises(ProxyException) as exc_info: + await call_cancel(cancel_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + cancel_harness.creds_resolver.assert_not_called() + cancel_harness.litellm_acancel.assert_not_called() diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5faae166fca..548c0eb0d91 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1877,7 +1877,7 @@ def test_get_file_content_streams_openai_direct_path( monkeypatch.setattr(litellm, "afile_content", _mock_afile_content) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: (False, None, None, None), + AsyncMock(return_value=(False, None, None, None)), ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -1942,15 +1942,17 @@ def test_get_file_content_routed_provider_skips_streaming_when_resolved_provider ) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: ( - True, - "azure-gpt-3-5-turbo", - "file-original-123", - { - "custom_llm_provider": "azure", - "api_key": "azure-key", - "api_base": "https://azure.example.com", - }, + AsyncMock( + return_value=( + True, + "azure-gpt-3-5-turbo", + "file-original-123", + { + "custom_llm_provider": "azure", + "api_key": "azure-key", + "api_base": "https://azure.example.com", + }, + ) ), ) @@ -2015,7 +2017,7 @@ def test_get_file_content_non_openai_provider_skips_streaming_handler( ) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: (False, None, None, None), + AsyncMock(return_value=(False, None, None, None)), ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -2463,14 +2465,16 @@ def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice( monkeypatch.setattr(litellm, "afile_list", _mock_afile_list) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: ( - True, - "azure-gpt-4o", - None, - { - "custom_llm_provider": "azure", - "api_key": "azure-key", - }, + AsyncMock( + return_value=( + True, + "azure-gpt-4o", + None, + { + "custom_llm_provider": "azure", + "api_key": "azure-key", + }, + ) ), ) @@ -4878,3 +4882,145 @@ def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFix assert captured_kwargs["api_key"] == "mistral-key" assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df" assert response.json()["id"] == encoded_id + + +def _mistral_plus_anthropic_router() -> Router: + return Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"}, + "model_info": {"id": "mistral-ocr-id"}, + }, + { + "model_name": "claude-opus-4-6", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "anthropic-key"}, + "model_info": {"id": "claude-id"}, + }, + ] + ) + + +def _restricted_key(key_models: list[str]) -> UserAPIKeyAuth: + from litellm.proxy._types import LitellmUserRoles + + return UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + team_id="team-a", + team_models=["claude-opus-4-6", "mistral-ocr"], + models=key_models, + ) + + +@pytest.mark.parametrize( + "http_method, path_suffix, litellm_fn", + [ + ("get", "", "afile_retrieve"), + ("get", "/content", "afile_content"), + ("delete", "", "afile_delete"), + ], +) +def test_model_routed_file_ops_reject_key_without_model_grant( + mocker: MockerFixture, monkeypatch, http_method: str, path_suffix: str, litellm_fn: str +): + """ + Regression: a key whose allowlist does not include the deployment named in a + model-encoded file id must be refused before that deployment's server-side + credentials are resolved. Previously any key could name any deployment via the + id (or the x-litellm-model header) and act on that provider account's files. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called")) + monkeypatch.setattr(litellm, litellm_fn, upstream) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"]) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = getattr(client, http_method)( + f"/v1/files/{encoded_id}{path_suffix}", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + assert "not allowed to access model" in response.text + upstream.assert_not_called() + + +def test_list_files_header_model_rejects_key_without_model_grant(mocker: MockerFixture, monkeypatch): + import litellm.proxy.proxy_server as ps + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called")) + monkeypatch.setattr(litellm, "afile_list", upstream) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"]) + + try: + response = client.get( + "/v1/files", headers={"Authorization": "Bearer test-key", "x-litellm-model": "mistral-ocr"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + upstream.assert_not_called() + + +def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFixture, monkeypatch): + """The grant check must not break the happy path: a key allowed the deployment still resolves its credentials.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_retrieve(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", + object="file", + bytes=2, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["mistral-ocr"]) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["api_key"] == "mistral-key" + assert captured_kwargs["custom_llm_provider"] == "mistral" diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index 3d1831bb4cd..fe4903b547c 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -58,10 +58,7 @@ def _make_batch_response( def test_get_batch_id_from_unified_batch_id_handles_appended_fields(): - decoded_id = ( - "litellm_proxy;model_id:deployment-123;" - "llm_batch_id:batch_openai_123;llm_output_file_id:file-output" - ) + decoded_id = "litellm_proxy;model_id:deployment-123;llm_batch_id:batch_openai_123;llm_output_file_id:file-output" assert get_batch_id_from_unified_batch_id(decoded_id) == "batch_openai_123" @@ -107,12 +104,10 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id(): } ), ), + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, + "litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model", + new=AsyncMock(return_value=mock_credentials), ), patch( "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", @@ -165,23 +160,15 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id(): ) # The batch_id should be encoded with model info - assert ( - response.id != raw_batch_id - ), f"Expected batch_id to be encoded, but got raw ID: {response.id}" - assert response.id.startswith( - "batch_" - ), f"Encoded batch_id should keep batch_ prefix, got: {response.id}" + assert response.id != raw_batch_id, f"Expected batch_id to be encoded, but got raw ID: {response.id}" + assert response.id.startswith("batch_"), f"Encoded batch_id should keep batch_ prefix, got: {response.id}" # Should be decodable back to the original decoded_model = decode_model_from_file_id(response.id) - assert ( - decoded_model == model_name - ), f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" + assert decoded_model == model_name, f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" original_id = get_original_file_id(response.id) - assert ( - original_id == raw_batch_id - ), f"Expected original ID '{raw_batch_id}', got: {original_id}" + assert original_id == raw_batch_id, f"Expected original ID '{raw_batch_id}', got: {original_id}" assert mock_create_batch.call_args.kwargs["metadata"] == {"customer_id": "cust-123"} @@ -227,12 +214,10 @@ async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_i } ), ), + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, + "litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model", + new=AsyncMock(return_value=mock_credentials), ), patch( "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", @@ -316,9 +301,7 @@ async def test_create_batch_without_x_litellm_model_returns_raw_ids(monkeypatch) } ), ), - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( "litellm.acreate_batch", new=AsyncMock(return_value=mock_response), @@ -383,9 +366,7 @@ class TestBatchIdRoundTripWithRetrieve: raw_batch_id = "batch_vllm_12345" # What create_batch does: - encoded_id = encode_file_id_with_model( - file_id=raw_batch_id, model=model_name, id_type="batch" - ) + encoded_id = encode_file_id_with_model(file_id=raw_batch_id, model=model_name, id_type="batch") # What retrieve_batch does: decoded_model = decode_model_from_file_id(encoded_id) @@ -410,9 +391,7 @@ class TestBatchIdRoundTripWithRetrieve: ] for raw_id, model in test_cases: - encoded = encode_file_id_with_model( - file_id=raw_id, model=model, id_type="batch" - ) + encoded = encode_file_id_with_model(file_id=raw_id, model=model, id_type="batch") assert encoded.startswith("batch_") assert decode_model_from_file_id(encoded) == model assert get_original_file_id(encoded) == raw_id @@ -440,9 +419,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ mock_user_api_key_dict.team_metadata = {} with ( - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( "litellm.proxy.batches_endpoints.endpoints.update_batch_in_database", new=AsyncMock(), From bae731ddfc394b23d3c6f44a85cfaa58472897e7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 17:54:57 -0400 Subject: [PATCH 020/349] fix(proxy): apply model grants to unified file and batch ids on batch routes Unified ids carry the deployment model inside the id, so a restricted key could create, retrieve or cancel a batch on a deployment it is not granted. The model parsed from a unified id now goes through the same grant check as header, query and model-encoded id sources before the router is called. --- litellm/proxy/batches_endpoints/endpoints.py | 16 ++++- .../proxy/batches_endpoints/test_endpoints.py | 58 +++++++++++++++++++ .../test_batch_x_litellm_model_encoding.py | 7 +-- 3 files changed, 75 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index c99f66d032e..c2489ce52ea 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, + authorize_model_for_key, batch_cost_poller_is_active, decode_model_from_file_id, encode_batch_response_ids, @@ -286,6 +287,7 @@ async def create_batch( detail={"error": f"Expected 1 model, got {len(target_model_names)}"}, ) model: Final = target_model_names[0] + await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict) _create_batch_data["model"] = model resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id) @@ -582,10 +584,17 @@ async def retrieve_batch( ) if unified_batch_id: + unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) + if unified_model_id is not None: + await authorize_model_for_key( + model_id=llm_router.resolve_model_name_from_model_id(unified_model_id) or unified_model_id, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) add_internal_model_credentials( data=data, llm_router=llm_router, - model_id=get_model_id_from_unified_batch_id(unified_batch_id), + model_id=unified_model_id, ) response = await llm_router.aretrieve_batch(**data) @@ -998,6 +1007,11 @@ async def cancel_batch( status_code=400, detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."}, ) + await authorize_model_for_key( + model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) data["model"] = model_id_from_batch data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id) response = await llm_router.acancel_batch(**data) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 57e42e79a42..5be64ca0847 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -179,6 +179,8 @@ def harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1165,6 +1167,8 @@ def retrieve_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1622,6 +1626,8 @@ def list_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2020,6 +2026,8 @@ def cancel_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2816,3 +2824,53 @@ async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_h assert exc_info.value.code == "403" cancel_harness.creds_resolver.assert_not_called() cancel_harness.litellm_acancel.assert_not_called() + + +def _b64_unified_id(decoded: str) -> str: + return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=") + + +UNIFIED_FILE_ID_FOR_GPT4O_MINI = _b64_unified_id( + "litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;" + "target_model_names,gpt-4o-mini;llm_output_file_id,file-provider;llm_output_file_model_id,dep-1" +) +UNIFIED_BATCH_ID_FOR_GPT4O_MINI = _b64_unified_id(UNIFIED_BATCH_ID) + + +@pytest.mark.asyncio +async def test_create__unified_file_id_rejects_key_without_model_grant(harness): + """The model carried inside a unified file id is caller-controlled too, so it is checked against the key's grants.""" + set_body( + harness, + { + "input_file_id": UNIFIED_FILE_ID_FOR_GPT4O_MINI, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_rejects_key_without_model_grant(retrieve_harness): + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + retrieve_harness.router_aretrieve.assert_not_called() + retrieve_harness.creds_resolver.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_harness): + with pytest.raises(ProxyException) as exc_info: + await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + cancel_harness.router_acancel.assert_not_called() diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index fe4903b547c..3161fe99e68 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, get_batch_id_from_unified_batch_id, @@ -412,11 +413,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ mock_request.url.path = f"/v1/batches/{unified_batch_id}/cancel" mock_fastapi_response = MagicMock() mock_fastapi_response.headers = {} - mock_user_api_key_dict = MagicMock() - mock_user_api_key_dict.parent_otel_span = None - mock_user_api_key_dict.user_id = "test_user" - mock_user_api_key_dict.allowed_model_region = None - mock_user_api_key_dict.team_metadata = {} + mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="test_user", team_metadata={}) with ( patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, From 95fdefa390af6586affb5ff825955b0ccb3bce17 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 19:14:45 -0400 Subject: [PATCH 021/349] fix(logging): tolerate a missing api_base in pre_call for presigned batch retrieves Provider batch configs that build their own request URL (Mistral, Bedrock) hand pre_call api_base=None, and mask_api_base_credentials raised TypeError on it, so every such retrieve logged a non-blocking LoggingError and lost its pre-call logging. --- litellm/litellm_core_utils/litellm_logging.py | 4 +- .../test_litellm_logging.py | 63 ++++++++++--------- 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cb9209be267..38e88493986 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1199,8 +1199,8 @@ class Logging(LiteLLMLoggingBaseClass): return {"error": f"Unable to parse raw request body. Got - {data}"} return data - def _get_masked_api_base(self, api_base: str) -> str: - return str(mask_api_base_credentials(api_base)) + def _get_masked_api_base(self, api_base: str | None) -> str: + return str(mask_api_base_credentials(api_base or "")) def _pre_call(self, input, api_key, model=None, additional_args={}): """ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 6aa77745e3d..570f4339cd4 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -58,6 +58,16 @@ def test_get_masked_api_base(logging_obj): assert type(masked_api_base) == str +def test_pre_call_tolerates_missing_api_base(logging_obj): + """Presigned batch retrieves (Mistral, Bedrock) build their own URL and pass api_base=None + to pre_call; masking must not raise or the request's pre-call logging is silently lost.""" + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="", additional_args={"api_base": None, "headers": {}}) + + assert logging_obj.model_call_details["litellm_params"]["api_base"] == "" + + def test_post_call_serializes_dict_with_datetime(logging_obj): import datetime @@ -3976,9 +3986,7 @@ def test_get_standard_logging_object_payload_carries_matched_access_groups(loggi "model": "gpt-4o", "messages": [], "litellm_params": { - "metadata": { - "user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"] - }, + "metadata": {"user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"]}, "proxy_server_request": {"body": {}}, }, }, @@ -4062,9 +4070,7 @@ def _model_router_response(selected_model: str, stamp: bool): from litellm.types.utils import ModelResponse response = ModelResponse(model=selected_model) - response._hidden_params = ( - {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} - ) + response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} return response @@ -4088,9 +4094,7 @@ def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response( - "azure_ai/grok-4-1-fast-reasoning", stamp=True - ), + init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True), start_time=now, end_time=now, logging_obj=logging_obj, @@ -4122,9 +4126,7 @@ def test_standard_logging_payload_keeps_requested_model_without_router_stamp( "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response( - "azure_ai/grok-4-1-fast-reasoning", stamp=False - ), + init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False), start_time=now, end_time=now, logging_obj=logging_obj, @@ -5536,9 +5538,7 @@ class TestNonInferenceCallTypesAreNotBilled: init_response_obj=self._retrieved_response(), start_time=now, end_time=now, - logging_obj=self._logging_obj( - "aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA - ), + logging_obj=self._logging_obj("aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA), status="success", ) @@ -5784,9 +5784,7 @@ async def test_streaming_success_callbacks_survive_cost_calculation_failure(): releasing.async_log_success_event = AsyncMock() patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) - with patcher, patch.object( - logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block") - ): + with patcher, patch.object(logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block")): await logging_obj.async_success_handler(result=_assembled_stream_result()) assert logging_obj.model_call_details["response_cost"] is None @@ -5799,8 +5797,9 @@ async def test_streaming_success_callbacks_survive_standard_logging_payload_fail releasing.async_log_success_event = AsyncMock() patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) - with patcher, patch.object( - logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream") + with ( + patcher, + patch.object(logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream")), ): await logging_obj.async_success_handler(result=_assembled_stream_result()) @@ -6073,6 +6072,8 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa ) for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook) + + def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch): """With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic" callback builds the OTel v2 logger (per-team credential routing); with the @@ -6228,7 +6229,9 @@ def test_get_error_information_skips_traceback_for_budget_rejection_with_provide from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup assert litellm.log_client_error_tracebacks is False - over_budget = _raise_and_catch(litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic")) + over_budget = _raise_and_catch( + litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic") + ) result = StandardLoggingPayloadSetup.get_error_information(over_budget) assert result["error_code"] == "429" assert result["llm_provider"] == "anthropic" @@ -6745,9 +6748,7 @@ def test_passthrough_embeddings_result_swapped_for_callbacks(): ], "model": "EmbeddingsGigaR", }, - request=httpx.Request( - "POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings" - ), + request=httpx.Request("POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings"), ) _, _, swapped_result = logging_obj._success_handler_helper_fn( @@ -6766,12 +6767,14 @@ def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervene request-level guardrail_status but never mask an intervention.""" flagged = {"guardrail_status": "guardrail_flagged"} - assert _get_status_fields( - "success", [{"guardrail_status": "success"}, flagged], None - )["guardrail_status"] == "guardrail_flagged" - assert _get_status_fields( - "success", [flagged, {"guardrail_status": "guardrail_intervened"}], None - )["guardrail_status"] == "guardrail_intervened" + assert ( + _get_status_fields("success", [{"guardrail_status": "success"}, flagged], None)["guardrail_status"] + == "guardrail_flagged" + ) + assert ( + _get_status_fields("success", [flagged, {"guardrail_status": "guardrail_intervened"}], None)["guardrail_status"] + == "guardrail_intervened" + ) def test_get_error_information_redacts_provider_key_from_upstream_url(): From e6bc4e47c7a63e8540a856f1b2f66710a4b47142 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 19:46:55 -0400 Subject: [PATCH 022/349] fix(mistral): reject file purposes Mistral lacks instead of mapping them to batch The proxy runs batch-file validation and guardrails only for purpose=batch, so a purpose such as assistants that was silently rewritten to batch on the way to Mistral let an upload skip both. Only batch, fine-tune and ocr pass through now; anything else is a 400. --- litellm/llms/mistral/files/transformation.py | 9 ++++-- .../test_mistral_files_transformation.py | 29 ++++++++++++++----- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 6d58311813c..bf1ef7cb69f 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -89,11 +89,14 @@ def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose: def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: + """Only Mistral's own purposes pass through. Silently mapping anything else to ``batch`` + would let an upload skip the proxy's batch-file validation and guardrails, which only + run when the caller says ``purpose=batch``.""" match purpose: - case "fine-tune" | "ocr": + case "batch" | "fine-tune" | "ocr": return purpose case _: - return "batch" + raise ValueError(f"Mistral does not support purpose={purpose!r}. Use one of: batch, fine-tune, ocr") def _api_base_from(litellm_params: Mapping[str, object]) -> str: @@ -166,7 +169,7 @@ class MistralFilesConfig(BaseFilesConfig): content_type: Final = extracted.get("content_type") or "application/octet-stream" upload: Final = MistralMultipartUpload( file=(filename, extracted["content"], content_type), - purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + purpose=(None, _to_mistral_purpose(create_file_data.get("purpose") or "batch")), ) return dict(upload) # mutable-ok: BaseFilesConfig signature diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index f62645be7ee..b81740c0429 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -91,18 +91,28 @@ def test_upload_request_is_multipart_with_batch_purpose(config): } -@pytest.mark.parametrize( - "openai_purpose,mistral_purpose", - [("batch", "batch"), ("fine-tune", "fine-tune"), ("ocr", "ocr"), ("assistants", "batch"), ("user_data", "batch")], -) -def test_upload_request_maps_purpose_onto_mistral_enum(config, openai_purpose, mistral_purpose): +@pytest.mark.parametrize("purpose", ["batch", "fine-tune", "ocr"]) +def test_upload_request_passes_mistral_purposes_through(config, purpose): body = config.transform_create_file_request( model="", - create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=openai_purpose), + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), optional_params={}, litellm_params={}, ) - assert body["purpose"] == (None, mistral_purpose) + assert body["purpose"] == (None, purpose) + + +@pytest.mark.parametrize("purpose", ["assistants", "user_data", "vision", "evals"]) +def test_upload_request_rejects_purposes_mistral_lacks(config, purpose): + """Regression: these used to be silently rewritten to ``batch``, so an upload that skipped the + proxy's batch-only validation and guardrails still landed on Mistral as a batch input file.""" + with pytest.raises(ValueError, match=f"purpose={purpose!r}"): + config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), + optional_params={}, + litellm_params={}, + ) def test_upload_request_requires_file(config): @@ -181,6 +191,11 @@ def test_list_request_filters_by_mapped_purpose(config): assert no_params == {} +def test_list_request_rejects_purposes_mistral_lacks(config): + with pytest.raises(ValueError, match="purpose='assistants'"): + config.transform_list_files_request(purpose="assistants", optional_params={}, litellm_params={}) + + def test_list_response(config): out = config.transform_list_files_response( raw_response=_response( From 0f3c4ccfbba894c0a66d92ed44ca18f312bf75f0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 11 Sep 2026 00:00:17 +0000 Subject: [PATCH 023/349] fix(auth): inherit organization_alias from the org for JWT and team-linked keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 42 ++++++- .../proxy/auth/test_user_api_key_auth.py | 104 +++++++++++++++++- 2 files changed, 143 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 20ab9904f46..110c524ecdf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( get_end_user_object, get_jwt_key_mapping_object, get_object_permission, + get_org_object, get_project_object, get_team_object, get_user_object, @@ -2398,6 +2399,37 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() +async def _inherit_org_identity( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTableCachedObj | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> None: + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + if ( + user_api_key_auth_obj.org_id is None + or user_api_key_auth_obj.organization_alias is not None + or prisma_client is None + ): + return + try: + org_object: Final = await get_org_object( + org_id=user_api_key_auth_obj.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.debug("org alias lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) + return + if org_object is not None: + user_api_key_auth_obj.organization_alias = org_object.organization_alias + + @tracer.wrap() async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, @@ -2622,8 +2654,14 @@ async def _run_centralized_common_checks( ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result - if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: - user_api_key_auth_obj.org_id = team_object.organization_id + await _inherit_org_identity( + user_api_key_auth_obj=user_api_key_auth_obj, + team_object=cast(LiteLLM_TeamTableCachedObj | None, team_object), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6cce6d0316b..03efbfa7185 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -31,7 +32,7 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object +from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, @@ -5293,6 +5294,107 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_id,team_org_id,existing_alias,lookup_mode,expected_org_id,expected_alias", + [ + (None, "t1", "org-from-team", None, "success", "org-from-team", "acme-org"), + ("org-jwt", None, None, None, "success", "org-jwt", "acme-org"), + ("org-pinned", None, None, "preset", "success", "org-pinned", "preset"), + ("org-missing", None, None, None, "missing", "org-missing", None), + ], +) +async def test_centralized_common_checks_inherits_org_alias( + key_org_id, + team_id, + team_org_id, + existing_alias, + lookup_mode, + expected_org_id, + expected_alias, +): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id=team_id, + org_id=key_org_id, + organization_alias=existing_alias, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = ( + LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) if team_id is not None else None + ) + organization = LiteLLM_OrganizationTable( + organization_id=expected_org_id, + organization_alias="acme-org", + budget_id="budget-id", + models=[], + created_by="test", + updated_by="test", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["prisma_client"] = MagicMock() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + identity_seen_by_common_checks = [] + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ) as mock_get_team_object, + patch( + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + return_value=organization, + ) as mock_get_org_object, + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + side_effect=lambda **kw: identity_seen_by_common_checks.append( + (kw["valid_token"].org_id, kw["valid_token"].organization_alias) + ), + ) as mock_checks, + ): + if lookup_mode == "missing": + mock_get_org_object.side_effect = OrganizationNotFoundError("x") + + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id == expected_org_id + assert token.organization_alias == expected_alias + assert identity_seen_by_common_checks == [(expected_org_id, expected_alias)] + if team_id is None: + mock_get_team_object.assert_not_awaited() + else: + mock_get_team_object.assert_awaited_once() + if existing_alias is not None: + mock_get_org_object.assert_not_awaited() + else: + mock_get_org_object.assert_awaited_once() + assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_cli_session_token_org_backfilled_from_team(monkeypatch): """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted From 94032014df175c3ec3735ded5144e91b5576547b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:24:59 -0700 Subject: [PATCH 024/349] fix(proxy): coordinate v2 migration startup and add container regression CI --- .circleci/config.yml | 137 +++++++++++- .circleci/scripts/run_migration_tests.py | 113 ++++++++++ .../litellm_proxy_extras/migration_lock.py | 89 ++++++++ .../migration_recovery.py | 158 ++++++++++++++ .../litellm_proxy_extras/prisma_toolchain.py | 5 + .../litellm_proxy_extras/utils.py | 192 +++++++++++------ .../tests/test_setup_database_fail_fast.py | 109 ++++------ tests/e2e/CLAUDE.md | 2 + tests/e2e/conftest.py | 10 +- tests/e2e/migrations/__init__.py | 0 tests/e2e/migrations/checks.py | 130 +++++++++++ tests/e2e/migrations/conftest.py | 62 ++++++ tests/e2e/migrations/containers.py | 203 ++++++++++++++++++ tests/e2e/migrations/database.py | 135 ++++++++++++ tests/e2e/migrations/startup_models.py | 25 +++ tests/e2e/migrations/test_legacy.py | 84 ++++++++ tests/e2e/migrations/test_pooling.py | 135 ++++++++++++ tests/e2e/migrations/test_recovery.py | 183 ++++++++++++++++ tests/e2e/migrations/test_startup.py | 100 +++++++++ .../test_litellm_proxy_extras_utils.py | 191 ++++++++++++---- .../test_migration_ci.py | 36 ++++ 21 files changed, 1909 insertions(+), 190 deletions(-) create mode 100644 .circleci/scripts/run_migration_tests.py create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_lock.py create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py create mode 100644 tests/e2e/migrations/__init__.py create mode 100644 tests/e2e/migrations/checks.py create mode 100644 tests/e2e/migrations/conftest.py create mode 100644 tests/e2e/migrations/containers.py create mode 100644 tests/e2e/migrations/database.py create mode 100644 tests/e2e/migrations/startup_models.py create mode 100644 tests/e2e/migrations/test_legacy.py create mode 100644 tests/e2e/migrations/test_pooling.py create mode 100644 tests/e2e/migrations/test_recovery.py create mode 100644 tests/e2e/migrations/test_startup.py create mode 100644 tests/proxy_migration_tests/test_migration_ci.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 32d2cf0390c..f6f31651306 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1,10 +1,32 @@ version: 2.1 +parameters: + run_migration_tests: + type: boolean + default: false + migration_candidate_image: + type: string + default: "" + migration_source_sha: + type: string + default: "" orbs: codecov: codecov/codecov@4.0.1 node: circleci/node@5.1.0 # Add this line to declare the node orb win: circleci/windows@5.0 # Add Windows orb commands: + checkout_migration_source: + steps: + - run: + name: Select the requested migration test revision + environment: + MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >> + command: | + if [ -n "$MIGRATION_SOURCE_SHA" ]; then + [[ "$MIGRATION_SOURCE_SHA" =~ ^[0-9a-f]{40}$ ]] || exit 1 + git fetch origin "$MIGRATION_SOURCE_SHA" + git checkout --detach "$MIGRATION_SOURCE_SHA" + fi skip_if_unrelated_changes: parameters: category: @@ -2853,14 +2875,25 @@ jobs: working_directory: ~/project steps: - checkout + - checkout_migration_source - skip_if_unrelated_changes - run: name: Build Docker image + environment: + MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >> command: | - docker build \ - -t litellm-docker-database:ci \ - -f docker/Dockerfile.database . + if [ -n "$MIGRATION_CANDIDATE_IMAGE" ]; then + [[ "$MIGRATION_CANDIDATE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+@sha256:[0-9a-f]{64}$ ]] || exit 1 + docker pull "$MIGRATION_CANDIDATE_IMAGE" + docker tag "$MIGRATION_CANDIDATE_IMAGE" litellm-docker-database:ci + else + docker build \ + --label org.opencontainers.image.revision="$(git rev-parse HEAD)" \ + -t litellm-docker-database:ci \ + -f docker/Dockerfile.database . + fi + python3 .circleci/scripts/run_migration_tests.py record-image - run: name: Save Docker image to workspace root @@ -2871,6 +2904,79 @@ jobs: root: . paths: - litellm-docker-database.tar.zst + - migration-image.json + + migration_startup_tests: + parameters: + suite: + type: enum + enum: [startup, recovery, legacy] + machine: + image: ubuntu-2204:2024.04.1 + resource_class: large + working_directory: ~/project + environment: + LITELLM_MIGRATION_TESTS: "1" + LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci + MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres + MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres + MIGRATION_TEST_OUTPUT: /tmp/migration-results + PYTHONPATH: tests/e2e + steps: + - checkout + - checkout_migration_source + - install_uv + - install_rust + - restore_cache: + keys: + - v1-uv-cache-{{ checksum "uv.lock" }} + - run: + name: Install test dependencies + command: uv sync --frozen --all-groups --all-extras --python 3.12 + - attach_workspace: + at: ~/project + - run: + name: Load the shared candidate and start PostgreSQL + command: | + zstd -d litellm-docker-database.tar.zst --stdout | docker load + docker run -d --name migration-postgres \ + -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres \ + -p 5432:5432 \ + postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" + - run: + name: Run migration startup regressions + environment: + MIGRATION_TEST_SUITE: << parameters.suite >> + MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >> + command: | + mkdir -p /tmp/migration-results + uv run --no-sync python .circleci/scripts/run_migration_tests.py + no_output_timeout: 15m + - store_test_results: + path: /tmp/migration-results/junit + - run: + name: Package migration diagnostics + when: always + command: | + mkdir -p /tmp/migration-artifacts + if [ -d /tmp/migration-results ]; then + tar -czf /tmp/migration-artifacts/diagnostics.tar.gz -C /tmp/migration-results . + if [ -f /tmp/migration-results/verdict.json ]; then + cp /tmp/migration-results/verdict.json /tmp/migration-artifacts/verdict.json + fi + fi + - store_artifacts: + path: /tmp/migration-artifacts + destination: migration-results + - run: + name: Remove migration test containers + when: always + command: | + docker ps -aq --filter label=litellm-migration-test=true | xargs -r docker rm -f + docker rm -f migration-postgres || true test_bad_database_url: machine: @@ -2915,7 +3021,32 @@ jobs: fi workflows: + migration_startup: + when: << pipeline.parameters.run_migration_tests >> + jobs: &migration_jobs + - build_docker_database_image + - migration_startup_tests: + name: migration-startup + suite: startup + requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-recovery + suite: recovery + requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-legacy-and-pooling + suite: legacy + requires: [build_docker_database_image] + migration_startup_scheduled: + triggers: + - schedule: + cron: "17 0,6,12,18 * * *" + filters: + branches: + only: litellm_internal_staging + jobs: *migration_jobs build_and_test: + unless: << pipeline.parameters.run_migration_tests >> jobs: - using_litellm_on_windows: filters: &main_branches diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py new file mode 100644 index 00000000000..56029c406fb --- /dev/null +++ b/.circleci/scripts/run_migration_tests.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +from pathlib import Path +from typing import Final +from xml.etree import ElementTree + +SUITES: Final = { + "startup": (("test_startup.py",), 12), + "recovery": (("test_recovery.py",), 15), + "legacy": (("test_legacy.py", "test_pooling.py"), 11), +} + + +def successful_junit(path: Path, expected: int, exit_code: int) -> bool: + if exit_code != 0 or not path.is_file(): + return False + try: + root: Final = ElementTree.parse(path).getroot() + except ElementTree.ParseError: + return False + cases: Final = tuple(root.iter("testcase")) + identities: Final = frozenset((case.get("classname"), case.get("name")) for case in cases) + return len(cases) == len(identities) == expected and all( + not any(case.find(tag) is not None for tag in ("failure", "error", "skipped")) for case in cases + ) + + +def output(*command: str) -> str: + return subprocess.check_output(command, text=True, timeout=90).strip() + + +def record_image() -> None: + source: Final = output("git", "rev-parse", "HEAD") + image: Final = output("docker", "image", "inspect", "litellm-docker-database:ci", "--format", "{{.Id}}") + revision: Final = output( + "docker", + "image", + "inspect", + "litellm-docker-database:ci", + "--format", + '{{index .Config.Labels "org.opencontainers.image.revision"}}', + ) + assert re.fullmatch(r"[0-9a-f]{40}", source), "Invalid source revision" + assert revision == source, "Candidate image revision differs from the tested source" + Path("migration-image.json").write_text( + json.dumps( + { + "source_sha": source, + "image_id": image, + "candidate_image": os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), + } + ) + ) + + +def main() -> int: + suite: Final = os.environ["MIGRATION_TEST_SUITE"] + files, expected = SUITES[suite] + metadata: Final = json.loads(Path("migration-image.json").read_text()) + assert metadata["source_sha"] == output("git", "rev-parse", "HEAD"), "Image and test source revisions differ" + assert metadata["image_id"] == output( + "docker", "image", "inspect", os.environ["LITELLM_MIGRATION_TEST_IMAGE"], "--format", "{{.Id}}" + ), "Loaded image differs from the build output" + assert metadata["candidate_image"] == os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), "Wrong release candidate" + destination: Final = Path(os.environ["MIGRATION_TEST_OUTPUT"]) + junit: Final = destination / "junit" / "results.xml" + junit.parent.mkdir(parents=True, exist_ok=True) + result: Final = subprocess.run( + ( + sys.executable, + "-m", + "pytest", + *(f"tests/e2e/migrations/{name}" for name in files), + "-vv", + "--tb=short", + "--durations=10", + f"--junitxml={junit}", + "-o", + "addopts=", + "--reruns=0", + ), + check=False, + timeout=1200, + ) + passed: Final = successful_junit(junit, expected, result.returncode) + (destination / "verdict.json").write_text( + json.dumps( + { + **metadata, + "suite": suite, + "expected_cases": expected, + "passed": passed, + "pytest_exit_code": result.returncode, + "test_revision": metadata["source_sha"], + "workflow_id": os.environ.get("CIRCLE_WORKFLOW_ID", ""), + "job_number": os.environ.get("CIRCLE_BUILD_NUM", ""), + }, + indent=2, + ) + ) + return 0 if passed else 1 + + +if __name__ == "__main__": + if sys.argv[1:] == ["record-image"]: + record_image() + else: + raise SystemExit(main()) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py new file mode 100644 index 00000000000..e4ccbe585a9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py @@ -0,0 +1,89 @@ +import random +import time +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +from litellm_proxy_extras._logging import logger +from litellm_proxy_extras.prisma_toolchain import MIGRATION_LOCK_TIMEOUT_ENV_VAR, migration_lock_timeout + +MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big") + +if TYPE_CHECKING: + import psycopg + + +def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]: + database_url: Final = environment.get("DATABASE_URL") + direct_url: Final = environment.get("DIRECT_URL") + if not database_url or not direct_url: + return environment + schema: Final = next((value for key, value in parse_qsl(urlsplit(database_url).query) if key == "schema"), "public") + direct: Final = urlsplit(direct_url) + parameters: Final = tuple((key, value) for key, value in parse_qsl(direct.query) if key != "schema") + return { + **environment, + "DATABASE_URL": urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema))))), + } + + +@dataclass(frozen=True, slots=True) +class _LockResult: + acquired: bool + + +def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = MIGRATION_LOCK_KEY) -> bool: + from psycopg.rows import class_row + + with connection.cursor(row_factory=class_row(_LockResult)) as cursor: + row: Final = cursor.execute("SELECT pg_try_advisory_xact_lock(%s) AS acquired", (key,)).fetchone() + return row is not None and row.acquired + + +@dataclass(frozen=True, slots=True) +class MigrationCoordinator: + connection: "psycopg.Connection[tuple[object, ...]]" + + def check_connection(self) -> None: + self.connection.execute("SELECT 1") + + def acquire_prisma_lock(self) -> None: + deadline: Final = time.monotonic() + migration_lock_timeout() + while time.monotonic() < deadline: + if _try_lock(self.connection, 72707369): + return + time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic()))) + raise RuntimeError( + "Timed out waiting for Prisma's lock to recover migration history. LiteLLM startup has stopped. " + "Another migration or a pooled database session may still hold the lock. Check the database lock holder. " + "When using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler." + ) + + +@contextmanager +def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]: + import psycopg + + wait_seconds: Final = migration_lock_timeout() + deadline: Final = time.monotonic() + wait_seconds + try: + with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection: + coordinator: Final = MigrationCoordinator(connection) + logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds) + while time.monotonic() < deadline: + with connection.transaction(): + if _try_lock(connection): + logger.info("Acquired the v2 migration coordinator lock") + + yield coordinator + coordinator.check_connection() + return + time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic()))) + except psycopg.Error as exc: + raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc + raise RuntimeError( + f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. " + f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}." + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py new file mode 100644 index 00000000000..9202317c776 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -0,0 +1,158 @@ +import hashlib +import subprocess +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Final +from uuid import uuid4 + +from litellm_proxy_extras import prisma_toolchain +from litellm_proxy_extras._logging import logger +from litellm_proxy_extras.migration_lock import MigrationCoordinator + +if TYPE_CHECKING: + import psycopg + + +@dataclass(frozen=True, slots=True) +class MigrationProgress: + checksum: str + applied_steps_count: int + logs: str + id: str = "" + finished: bool = False + + def confirms_completion(self, script: bytes) -> bool: + return ( + self.applied_steps_count == 1 + and not self.logs.strip() + and self.checksum == hashlib.sha256(script).hexdigest() + ) + + +def _migration_records( + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, migration: Path +) -> tuple[MigrationProgress, ...]: + from psycopg import sql + from psycopg.rows import class_row + + with connection.cursor(row_factory=class_row(MigrationProgress)) as cursor: + records: Final = cursor.execute( + sql.SQL( + "SELECT id, checksum, applied_steps_count, coalesce(logs, '') AS logs, " + "finished_at IS NOT NULL AS finished FROM {} " + "WHERE migration_name = %s AND rolled_back_at IS NULL" + ).format(sql.Identifier(schema, "_prisma_migrations")), + (migration.parent.name,), + ).fetchall() + return tuple(records) + + +def recover_completed_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool: + """Finish a proven successful row without erasing its durable completion evidence. + + The caller commits this checkpoint before running another Prisma command. + """ + from psycopg import sql + + coordinator.acquire_prisma_lock() + records: Final = _migration_records(coordinator.connection, schema, migration) + unfinished: Final = tuple(record for record in records if not record.finished) + script: Final = migration.read_bytes() + if not unfinished: + return any(record.checksum == hashlib.sha256(script).hexdigest() for record in records) + if len(unfinished) != 1 or not unfinished[0].confirms_completion(script): + return False + progress: Final = unfinished[0] + result: Final = coordinator.connection.execute( + sql.SQL( + "UPDATE {} SET finished_at = current_timestamp " + "WHERE id = %s AND checksum = %s AND applied_steps_count = 1 " + "AND finished_at IS NULL AND rolled_back_at IS NULL AND coalesce(logs, '') = %s" + ).format(sql.Identifier(schema, "_prisma_migrations")), + (progress.id, progress.checksum, progress.logs), + ) + if result.rowcount != 1: + raise RuntimeError("Could not complete the confirmed migration history row; retry startup.") + logger.info("Completed migration %s using its successful SQL step and matching checksum", migration.parent.name) + return True + + +def migration_files(directory: Path) -> tuple[tuple[str, str], ...]: + return tuple( + (path.parent.name, hashlib.sha256(path.read_bytes()).hexdigest()) + for path in sorted((directory / "migrations").glob("*/migration.sql")) + ) + + +def baseline_current_schema( + coordinator: MigrationCoordinator, + schema: str, + migrations_dir: Path, + prisma_command: str, + prisma_env: Mapping[str, str], +) -> None: + from psycopg import sql + + packaged_dir: Final = Path(__file__).parent + migrations: Final = migration_files(migrations_dir) + if ( + not migrations + or migrations != migration_files(packaged_dir) + or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes() + ): + raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.") + + coordinator.acquire_prisma_lock() + existing: Final = coordinator.connection.execute( + "SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(coordinator.connection),) + ).fetchone() + if existing is not None and existing[0] is not None: + return + try: + prisma_toolchain.run_prisma( + ( + prisma_command, + "migrate", + "diff", + "--from-schema-datasource", + str(migrations_dir / "schema.prisma"), + "--to-schema-datamodel", + str(migrations_dir / "schema.prisma"), + "--exit-code", + ), + timeout=prisma_toolchain.prisma_command_timeout(), + env=prisma_env, + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc: + raise RuntimeError( + "Cannot automatically baseline this database: its schema has not been verified to match this build. " + "Establish the existing migration history before retrying. No schema reconciliation was performed. " + "If using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler. " + f"Schema verification detail: {exc.stderr}" + ) from exc + + coordinator.check_connection() + ledger: Final = sql.Identifier(schema, "_prisma_migrations") + coordinator.connection.execute( + sql.SQL( + "CREATE TABLE {} (id varchar(36) PRIMARY KEY NOT NULL, checksum varchar(64) NOT NULL, " + "finished_at timestamptz, migration_name varchar(255) NOT NULL, logs text, rolled_back_at timestamptz, " + "started_at timestamptz NOT NULL DEFAULT now(), applied_steps_count integer NOT NULL DEFAULT 0)" + ).format(ledger) + ) + with coordinator.connection.cursor() as cursor: + cursor.executemany( + sql.SQL( + "INSERT INTO {} (id, checksum, migration_name, logs, started_at, finished_at) " + "VALUES (%s, %s, %s, '', current_timestamp, current_timestamp)" + ).format(ledger), + tuple((str(uuid4()), checksum, name) for name, checksum in migrations), + ) + logger.warning( + "Legacy migration history was missing. The existing Prisma schema matches this build; " + "adopted %s packaged migrations as a baseline. No schema changes were applied, and " + "historical data backfills were not replayed or verified. Continuing startup; " + "review any feature-specific backfill requirements.", + len(migrations), + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 9cd48fcf11a..07f83f76d2b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -59,6 +59,7 @@ except ImportError: PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT" PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT" PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT" +MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT" NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR" DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0 @@ -106,6 +107,10 @@ def prisma_command_timeout() -> float: ) +def migration_lock_timeout() -> float: + return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0) + + def prisma_bootstrap_timeout() -> float: """Seconds the one-time Node toolchain install may run for.""" return _timeout_from_env( diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2145f891318..2749db5d754 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -6,6 +6,7 @@ import shutil import subprocess import tempfile import time +from collections.abc import Callable from dataclasses import dataclass, replace from pathlib import Path from typing import TYPE_CHECKING, Final, Optional @@ -78,15 +79,10 @@ MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 @dataclass(frozen=True) class _MigrateAttemptBudget: - """Retries left, and the recoveries already run. - - A recovery that lands something new costs nothing, so a database full of - objects `prisma db push` created works through them one per pass. Anything - that made no progress spends an attempt, so a stuck run still gives up. - """ + """Independent bounds for failed attempts and Prisma lock contention.""" attempts_left: int - recoveries: frozenset[str] = frozenset() + contention_seconds_left: float = 600.0 @property def exhausted(self) -> bool: @@ -99,10 +95,14 @@ class _MigrateAttemptBudget: def spend(self) -> "_MigrateAttemptBudget": return replace(self, attempts_left=self.attempts_left - 1) - def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget": - if recovery in self.recoveries: - return self.spend() - return replace(self, recoveries=self.recoveries | {recovery}) + def after_contention(self, elapsed: float) -> "_MigrateAttemptBudget": + remaining: Final = self.contention_seconds_left - elapsed + if remaining <= 0: + raise RuntimeError( + "Timed out waiting for Prisma's migration advisory lock. Check the running migration " + "or increase LITELLM_MIGRATION_LOCK_TIMEOUT." + ) + return replace(self, contention_seconds_left=remaining) _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) @@ -836,12 +836,47 @@ class ProxyExtrasDBManager: @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: + if not use_migrate: + return ProxyExtrasDBManager._run_database_v2(False) + from litellm_proxy_extras.migration_lock import migration_environment, migration_lock + from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration + + database_url: Final = os.environ.get("DATABASE_URL") + if not database_url: + raise RuntimeError("DATABASE_URL is required for v2 migrations") + lock_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.environ.get("DIRECT_URL") or database_url) + schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" + + def recover_completed(name: str) -> bool: + if Path(name).name != name or "\\" in name: + return False + migration: Final = Path(os.getcwd()) / "migrations" / name / "migration.sql" + if not migration.is_file(): + return False + with migration_lock(lock_url) as coordinator: + return recover_completed_migration(coordinator, schema, migration) + + def baseline_existing(migrations_dir: str) -> None: + with migration_lock(lock_url) as coordinator: + baseline_current_schema( + coordinator, schema, Path(migrations_dir), _get_prisma_command(), migration_environment(_get_prisma_env()) + ) + + while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing): + continue + return True + + @staticmethod + def _run_database_v2( + use_migrate: bool, + recover_completed: Callable[[str], bool] = lambda name: False, + baseline_existing: "Callable[[str], None] | None" = None, + ) -> bool: """ v2 migration resolver (opt-in via --use_v2_migration_resolver). - Runs `prisma migrate deploy` and handles standard recovery paths - (P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a - concurrent migrate deploy). Critically, it does + Runs `prisma migrate deploy`, baselines verified existing schemas, + and recovers confirmed SQL completion or reported deadlocks. It does NOT call `_resolve_all_migrations` — the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contended for the same DB during rolling deploys. @@ -850,10 +885,9 @@ class ProxyExtrasDBManager: is logged as a warning, not a fatal error — users whose DBs got into weird shapes from the old thrashing should still be able to start. - The retry budget only counts attempts that made no progress: see - _MigrateAttemptBudget. + False requests a committed recovery checkpoint and another deploy + pass. True means every pending migration is complete. """ - schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" migrations_dir = ProxyExtrasDBManager._get_prisma_dir() if not use_migrate: @@ -886,14 +920,22 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) deploy_timeout = prisma_migrate_deploy_timeout() - budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS) + from litellm_proxy_extras.migration_lock import migration_environment, migration_lock_timeout + + migration_env: Final = migration_environment(_get_prisma_env()) + + budget = _MigrateAttemptBudget( + attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS, + contention_seconds_left=migration_lock_timeout(), + ) try: while not budget.exhausted: + attempt_started = time.monotonic() try: result = prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "deploy"], timeout=deploy_timeout, - env=_get_prisma_env(), + env=migration_env, ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") return True @@ -909,8 +951,16 @@ class ProxyExtrasDBManager: next_budget = budget.spend() except subprocess.CalledProcessError as e: + if "P3005" in (e.stderr or "") and baseline_existing is not None: + baseline_existing(migrations_dir) + return False + failed_migration = ProxyExtrasDBManager._v2_failed_migration_name(e.stderr or "") + if failed_migration and recover_completed(failed_migration): + return False next_budget = ProxyExtrasDBManager._budget_after_deploy_failure( - e, budget, schema_path + e, + budget, + time.monotonic() - attempt_started, ) if next_budget.attempts_left < budget.attempts_left: @@ -919,19 +969,41 @@ class ProxyExtrasDBManager: raise RuntimeError( f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} " - "attempts that made no progress (timeouts, deadlock retries, or a " - "recovery that had already run once). Check database connectivity, " + "attempts that made no progress (timeouts or deadlock retries). Check database connectivity, " "load, and _prisma_migrations ledger state, and raise " f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out." ) finally: os.chdir(original_dir) + @staticmethod + def _v2_failed_migration_name(stderr: str) -> "str | None": + if "P3009" in stderr: + match = re.search(r"`(\d+_[^`\r\n]+)`", stderr) + return match.group(1) if match else None + if "P3018" in stderr: + match = re.search(r"Migration name: (\d+_[^\r\n]+)", stderr) + return match.group(1) if match else None + return None + + @staticmethod + def _v2_roll_back_migration_best_effort(migration_name: str) -> None: + from litellm_proxy_extras.migration_lock import migration_environment + + try: + prisma_toolchain.run_prisma( + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + timeout=prisma_command_timeout(), + env=migration_environment(_get_prisma_env()), + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass + @staticmethod def _budget_after_deploy_failure( error: subprocess.CalledProcessError, budget: "_MigrateAttemptBudget", - schema_path: str, + attempt_seconds: float = 0.0, ) -> "_MigrateAttemptBudget": """Recover from one failed `prisma migrate deploy`, and price the pass. @@ -940,37 +1012,35 @@ class ProxyExtrasDBManager: """ stderr = error.stderr or "" - if "P3005" in stderr and "database schema is not empty" in stderr: - logger.info("Schema exists but no migrations ledger — creating baseline") - if ProxyExtrasDBManager._create_baseline_migration(schema_path): - return budget.after_recovery("baseline") - return budget.spend() - if "P3009" in stderr: - migration_match = re.search(r"`(\d+_\S+?)`", stderr) - if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): - name = migration_match.group(1) - logger.info( - f"Migration {name} failed idempotently — marking applied and retrying" - ) - ProxyExtrasDBManager._mark_migration_applied(name) - return budget.after_recovery(f"resolved:{name}") - if migration_match: - migration_name = migration_match.group(1) + migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) + if migration_name: ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) - if ledger_logs is not None and ( - ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs - ): + if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs: logger.info( "Migration %s failed in a concurrent migrate deploy " "deadlock race, rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) return budget.spend() raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" + "Migration completion could not be verified. LiteLLM startup has stopped.\n\n" + f"Prisma migration history (migration name and start time):\n{stderr}\n\n" + "A migration has a start record but no successful completion record. " + "LiteLLM cannot determine whether its SQL committed from this record alone. " + "Startup stopped to avoid repeating or skipping database changes.\n\n" + "Before resolving, stop other migration runners and inspect _prisma_migrations, " + "the named migration.sql from this build, database logs, and the actual database objects and data. " + "Use the same database and this build's schema and migration files for recovery:\n" + "- Only after verifying every migration change is present, run " + "prisma migrate resolve --applied , then retry startup.\n" + "- Only after verifying no migration changes remain (or fully undoing partial changes), run " + "prisma migrate resolve --rolled-back , then retry startup. " + "This command updates history; it does not undo SQL.\n" + "Replace with the reported name. If the outcome remains uncertain, " + "leave migration history unchanged and contact your database administrator. " + "Repeated restarts alone will not resolve this state." ) from error if "P3018" in stderr: @@ -981,25 +1051,13 @@ class ProxyExtrasDBManager: f"and retry.\n\nPrisma error:\n{stderr}" ) from error - migration_match = re.search(r"Migration name: (\d+_\S+)", stderr) - if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): - name = migration_match.group(1) + migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) + if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr: logger.info( - f"Migration {name} SQL hit idempotent error — marking applied and retrying" - ) - ProxyExtrasDBManager._mark_migration_applied(name) - return budget.after_recovery(f"resolved:{name}") - - if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: - logger.info( - "Migration %s deadlocked against a concurrent " - "migrate deploy, rolling its ledger row back " - "and retrying", - migration_match.group(1), - ) - ProxyExtrasDBManager._roll_back_migration_best_effort( - migration_match.group(1) + "Migration %s deadlocked against a concurrent migrate deploy, rolling its ledger row back and retrying", + migration_name, ) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) return budget.spend() raise RuntimeError( @@ -1009,19 +1067,17 @@ class ProxyExtrasDBManager: if _MIGRATION_DEADLOCK_MARKER in stderr: logger.info( - "prisma migrate deploy attempt %s deadlocked against " - "a concurrent migrate deploy, retrying", + "prisma migrate deploy attempt %s deadlocked against a concurrent migrate deploy, retrying", budget.attempt_number, ) return budget.spend() if "P1002" in stderr and "advisory lock" in stderr: logger.info( - "prisma migrate deploy attempt %s timed out waiting for " - "the advisory lock a concurrent migrate deploy holds, retrying", - budget.attempt_number, + "Waiting for the advisory lock held by another Prisma migration; " + "contention does not spend a migration failure attempt" ) - return budget.spend() + return budget.after_contention(attempt_seconds) raise RuntimeError( "Database migration failed and cannot be auto-recovered. " diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 040d67d25e4..338c571eb4f 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -32,9 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str): def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): """v2: a permission failure during migrate deploy raises RuntimeError.""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -50,9 +48,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): """v2: a non-idempotent migration failure raises (no silent recovery).""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -61,7 +57,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): 'Reason: syntax error at or near "BRKN" LINE 42' ) with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -176,51 +172,33 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) -def test_v2_resolve_specific_migration_failure_raises_runtime_error( - monkeypatch, tmp_path -): - """If marking a migration as applied fails inside P3009 idempotent - recovery, the subprocess error must be re-raised as RuntimeError so - proxy_cli.py catches it cleanly (instead of leaking CalledProcessError).""" +def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path): + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "relation already exists") monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ProxyExtrasDBManager, + "_v2_roll_back_migration_best_effort", + lambda name: pytest.fail("duplicate-object errors do not prove rollback is safe"), ) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") monkeypatch.setattr( - ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("duplicate-object errors do not prove all SQL completed"), ) - - # First call: migrate deploy -> P3009 idempotent error. - # Recovery path tries _resolve_specific_migration; that also raises. - def _failing_resolve(*a, **kw): - raise subprocess.CalledProcessError( - returncode=1, - cmd="prisma migrate resolve --applied", - stderr="resolve failed", - output="", - ) - - monkeypatch.setattr( - ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve - ) - - stderr = ( - "Error: P3009\nMigration `20260101000000_some_migration` failed\n" - "relation already exists" - ) - with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises( - RuntimeError, match="Failed to mark migration .* as applied" - ): + stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists" + with patch( + "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) + ) as run: + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], + ) def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -252,9 +230,7 @@ _DEADLOCK_P3018_STDERR = ( def _stub_v2_env(monkeypatch, tmp_path): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") monkeypatch.setattr("time.sleep", lambda _: None) @@ -272,9 +248,7 @@ def _succeed_after(failures: int, stderr: str): return _OkResult() calls["n"] += 1 if calls["n"] <= failures: - raise subprocess.CalledProcessError( - returncode=1, cmd=args[0], stderr=stderr, output="" - ) + raise subprocess.CalledProcessError(returncode=1, cmd=args[0], stderr=stderr, output="") return _OkResult() return _run @@ -288,7 +262,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): rolled_back = [] monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", + "_v2_roll_back_migration_best_effort", lambda name: rolled_back.append(name), ) monkeypatch.setattr( @@ -306,7 +280,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): """v2: a deadlock on every attempt still fails after the retry budget.""" _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_v2_roll_back_migration_best_effort", lambda name: None) with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", @@ -335,7 +309,7 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ rolled_back = [] monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", + "_v2_roll_back_migration_best_effort", lambda name: rolled_back.append(name), ) monkeypatch.setattr( @@ -350,10 +324,8 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] -def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path): - """v2: empty failed ledger logs mean a concurrent deploy moved it on.""" +def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path): _stub_v2_env(monkeypatch, tmp_path) - stderr = ( "Error: P3009\n" "migrate found failed migrations in the target database\n" @@ -361,22 +333,19 @@ def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path "started at 2026-09-01 18:46:13 UTC failed" ) monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "") - rolled_back = [] monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", - lambda name: rolled_back.append(name), + "_v2_roll_back_migration_best_effort", + lambda name: pytest.fail("empty logs do not prove rollback is safe"), ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + with patch( + "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) + ) as run: + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], ) - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) - - ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) - assert ok is True - assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): @@ -392,12 +361,12 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None) monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", + "_v2_roll_back_migration_best_effort", lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), ) monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) - with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -418,7 +387,7 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): ) with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index a58c13d6a1c..9ef7cacf1a8 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -6,6 +6,8 @@ Code-style rules for writing tests under `tests/e2e/`. The harness already encod Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family or behavior area. If you add a new folder, you must add a line here describing what kind of tests belong in it, so the layout stays self-describing. `gateway/` is the exception: it holds proxy configuration only and never tests +- `migrations/` - isolated Docker startup, concurrent migration, crash recovery, and legacy database compatibility. The CircleCI migration workflow enables `LITELLM_MIGRATION_TESTS=1`; these tests own their proxy containers and databases, so they do not use the shared proxy preflight or shared database cleanup + - `llm_translation/` - LLM endpoint and provider-translation behavior: passthrough, custom pricing, OCR, and the non-chat inference endpoints (`/v1/responses`, `/v1/messages`, `/embeddings`, `/v1/rerank`, `/v1/audio/speech`, `/v1/images/generations`), each against a deployment the test creates via `/model/new` and deletes on teardown - `access_control/` - the gateway's authorization and error-shape contract: per-key model allow-lists, route-group permissions (`allowed_routes`), and unknown-model validation - `embeddings/` - the `/embeddings` endpoint across providers diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 36569896125..4a5f0aa880f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -64,6 +64,7 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line("markers", "migration_startup: isolated container startup tests run by the migration CI workflow") config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", @@ -123,6 +124,11 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: traffic only after the latency-sensitive suites have finished.""" for item in items: attach_result_properties(item) + if os.environ.get("LITELLM_MIGRATION_TESTS") != "1": + deselected = [item for item in items if item.get_closest_marker("migration_startup") is not None] + items[:] = [item for item in items if item.get_closest_marker("migration_startup") is None] + if deselected: + deselected[0].config.hook.pytest_deselected(items=deselected) items.sort(key=lambda item: item.get_closest_marker("load") is not None) @@ -155,7 +161,7 @@ def pytest_runtest_setup(item: pytest.Item) -> None: Unmarked tests (unit coverage of the harness) don't touch the proxy, so they run even when none is up. Never skip for a missing proxy. Replay mode needs the proxy too: only provider-bound traffic replays from the bundle.""" - if item.get_closest_marker("e2e") is None: + if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None: return reason = _proxy_fail_reason() if reason is not None: @@ -168,7 +174,7 @@ def pytest_runtest_call(item: pytest.Item) -> None: guard before truncating the spend-log DB. Tests under `tests/e2e/` without the `e2e` marker (pure unit coverage for the harness itself) never hit the proxy, so they must not arm the destructive DB truncate.""" - if item.get_closest_marker("e2e") is None: + if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None: return item.session.stash[_E2E_TEST_RAN] = True diff --git a/tests/e2e/migrations/__init__.py b/tests/e2e/migrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py new file mode 100644 index 00000000000..b14631ccad8 --- /dev/null +++ b/tests/e2e/migrations/checks.py @@ -0,0 +1,130 @@ +import hashlib +from contextlib import ExitStack +from typing import Final +from uuid import uuid4 + +from psycopg import sql + +from .containers import Containers, Replica, failed, until +from .database import GATE_KEY, Database +from .startup_models import Migration + +COMPLETE_SQL: Final = "CREATE TABLE migration_effect (id int PRIMARY KEY); INSERT INTO migration_effect VALUES (1);" +COMPLETE: Final = Migration("20990101000000_startup_test", COMPLETE_SQL) +NEXT: Final = Migration( + "20990102000000_next_test", + "CREATE TABLE migration_next (id int PRIMARY KEY); INSERT INTO migration_next VALUES (2);", +) +FATAL: Final = Migration(COMPLETE.name, "DO $$ BEGIN RAISE EXCEPTION 'MIGRATION_TEST_FATAL'; END $$;") +GATED: Final = Migration( + COMPLETE.name, f"SELECT pg_advisory_lock({GATE_KEY}); {COMPLETE.script} SELECT pg_advisory_unlock({GATE_KEY});" +) + + +def start_replicas( + stack: ExitStack, containers: Containers, database: Database, migrations: tuple[Migration, ...] = (), count: int = 3 +) -> tuple[Replica, ...]: + return tuple(stack.enter_context(containers.start(database, migrations)) for _ in range(count)) + + +def assert_completed(database: Database, migration: Migration = COMPLETE) -> None: + assert database.query( + "SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + (migration.name,), + ) == ((True, True, 1),), "Expected exactly one successful SQL execution" + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + +def confirmed_history(database: Database) -> str: + database.execute(COMPLETE_SQL) + row_id: Final = str(uuid4()) + database.execute( + "INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) VALUES (%s, %s, %s, 1)", + (row_id, COMPLETE.name, hashlib.sha256(COMPLETE.script.encode()).hexdigest()), + ) + return row_id + + +def assert_original_proof(database: Database, row_id: str, finished: bool) -> None: + assert database.query( + "SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", + (COMPLETE.name,), + ) == ((row_id, 1, finished, True),), "Recovery lost or replaced the original durable SQL proof" + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + +def pause_completion(database: Database) -> None: + database.execute( + sql.SQL( + "CREATE FUNCTION migration_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " + "IF NEW.migration_name = {name} AND NEW.finished_at IS NOT NULL THEN " + "PERFORM pg_advisory_lock({gate}); PERFORM pg_advisory_unlock({gate}); END IF; RETURN NEW; END $$; " + "CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW EXECUTE FUNCTION migration_pause()" + ).format(name=sql.Literal(COMPLETE.name), gate=sql.Literal(GATE_KEY)) + ) + + +def interrupt_owner( + containers: Containers, database: Database, after_commit: bool, *, stop_database_session: bool = True +) -> None: + if after_commit: + pause_completion(database) + with database.lock(): + with containers.start(database, (COMPLETE if after_commit else GATED,)) as owner: + until("migration at the intended crash boundary", lambda: bool(database.blocked())) + assert database.exists("migration_effect") == after_commit + assert database.query( + "SELECT finished_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,) + ) == ((True,),) + blocked: Final = database.blocked() + assert len(blocked) == 1 + backend: Final = blocked[0][0] + assert owner.state().Running + owner.kill() + assert owner.state().ExitCode == 137 + if stop_database_session: + database.query("SELECT pg_terminate_backend(%s)", (backend,)) + until( + "terminated migration backend released", + lambda: not database.query("SELECT pid FROM pg_stat_activity WHERE pid = %s", (backend,)), + ) + assert database.query( + "SELECT finished_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + (COMPLETE.name,), + ) == ((True, int(after_commit)),) + assert database.exists("migration_effect") == after_commit + if not stop_database_session: + until( + "database backend noticed container death", + lambda: not database.query("SELECT pid FROM pg_stat_activity WHERE pid = %s", (backend,)), + 60, + ) + + +def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None: + failed(replicas, "Migration completion could not be verified") + started: Final = str( + database.query( + "SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM _prisma_migrations WHERE migration_name = %s", + (COMPLETE.name,), + )[0][0] + ) + for replica in replicas: + assert_guidance(replica.logs(), started) + + +def assert_guidance(log: str, started: str) -> None: + for detail in ( + COMPLETE.name, + started, + "cannot determine whether its SQL committed", + "_prisma_migrations", + "migration.sql", + "Only after verifying every migration change is present", + "prisma migrate resolve --applied ", + "Only after verifying no migration changes remain", + "prisma migrate resolve --rolled-back ", + "leave migration history unchanged", + "Repeated restarts alone", + ): + assert detail in log, f"Missing recovery guidance: {detail}" diff --git a/tests/e2e/migrations/conftest.py b/tests/e2e/migrations/conftest.py new file mode 100644 index 00000000000..735adeedbdb --- /dev/null +++ b/tests/e2e/migrations/conftest.py @@ -0,0 +1,62 @@ +import json +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import pytest +from _pytest.fixtures import SubRequest + +from .containers import Containers, docker, ready +from .database import Database, Databases + + +@pytest.fixture(scope="session") +def migration_image(tmp_path_factory: pytest.TempPathFactory) -> str: + configured: Final = os.environ.get("LITELLM_MIGRATION_TEST_IMAGE") + assert configured, "LITELLM_MIGRATION_TEST_IMAGE must name the built candidate image" + image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}") + assert image.startswith("sha256:"), "Unable to identify the candidate image" + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) + output.mkdir(parents=True, exist_ok=True) + (output / "image.json").write_text(json.dumps({"requested": configured, "image_id": image})) + return image + + +@pytest.fixture(scope="session") +def databases() -> Databases: + admin: Final = os.environ.get("MIGRATION_TEST_ADMIN_URL", "") + parsed: Final = urlsplit(admin) + assert parsed.hostname in ("127.0.0.1", "localhost"), "Use an isolated loopback PostgreSQL test cluster" + assert parsed.port and parsed.path and not parsed.query, "Supply the test cluster port and admin database" + container_admin: Final = os.environ.get( + "MIGRATION_TEST_CONTAINER_ADMIN_URL", + admin.replace("127.0.0.1", "host.docker.internal").replace("localhost", "host.docker.internal"), + ) + return Databases(admin, container_admin) + + +@pytest.fixture(scope="session") +def migrated_template( + databases: Databases, migration_image: str, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Database]: + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "seed" + with databases.create() as database: + with Containers(migration_image, output).start(database) as replica: + ready((replica,), database) + yield database + + +@pytest.fixture +def database(databases: Databases, migrated_template: Database) -> Iterator[Database]: + with databases.create(migrated_template) as database: + yield database + + +@pytest.fixture +def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Containers: + configured: Final = os.environ.get("MIGRATION_TEST_OUTPUT") + output: Final = Path(configured) / request.node.name if configured else tmp_path + output.mkdir(parents=True, exist_ok=True) + return Containers(migration_image, output) diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py new file mode 100644 index 00000000000..0f5793b81dd --- /dev/null +++ b/tests/e2e/migrations/containers.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import hashlib +import subprocess +import time +from collections.abc import Callable, Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from uuid import uuid4 + +from e2e_http import NoBody, Success, unwrap +from models import KeyGenerateBody, KeyGenerateResponse, KeyInfoParams, KeyInfoResponse +from transport import HttpTransport + +from .database import Database, prisma_url +from .startup_models import ContainerState, Migration, Observation, Readiness + +MASTER_KEY: Final = "sk-migration-ci-fixture" + + +def docker(*args: str) -> str: + result: Final = subprocess.run(("docker", *args), capture_output=True, text=True, timeout=90) + assert result.returncode == 0, f"Docker operation failed: {result.stderr}" + return result.stdout.strip() + + +def until(description: str, condition: Callable[[], bool], seconds: float = 150) -> None: + deadline: Final = time.monotonic() + seconds + while time.monotonic() < deadline: + if condition(): + return + time.sleep(0.25) + raise AssertionError(f"Timed out waiting for {description}") + + +@dataclass(frozen=True, slots=True) +class Replica: + name: str + transport: HttpTransport + output: Path + + def state(self) -> ContainerState: + return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name)) + + def observe(self) -> Observation: + state: Final = self.state() + result: Final = self.transport.get( + "/health/readiness", headers=self.transport.master, params=NoBody(), response_type=Readiness, timeout=1 + ) + ready: Final = isinstance(result, Success) and result.data.status == "healthy" and result.data.db == "connected" + return Observation(None if state.Running else state.ExitCode, ready) + + def logs(self) -> str: + result: Final = subprocess.run(("docker", "logs", self.name), capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stderr + return result.stdout + result.stderr + + def kill(self) -> None: + if self.state().Running: + docker("kill", self.name) + + def usable(self, database: Database) -> None: + alias: Final = f"migration-{uuid4().hex}" + key: Final = unwrap( + self.transport.post( + "/key/generate", + headers=self.transport.master, + json=KeyGenerateBody(key_alias=alias), + response_type=KeyGenerateResponse, + ) + ).key + info: Final = unwrap( + self.transport.get( + "/key/info", + headers=self.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ) + assert info.info.key_alias == alias + assert database.query( + 'SELECT key_alias FROM "LiteLLM_VerificationToken" WHERE token = %s', + (hashlib.sha256(key.encode()).hexdigest(),), + ) == ((alias,),) + + +def ready(replicas: tuple[Replica, ...], database: Database) -> None: + def all_ready() -> bool: + observations: Final = tuple(replica.observe() for replica in replicas) + assert all(item.exit_code is None for item in observations), "Replica exited before readiness" + return all(item.ready for item in observations) + + until("every replica ready", all_ready) + for replica in replicas: + replica.usable(database) + + +def failed(replicas: tuple[Replica, ...], marker: str) -> None: + def all_stopped() -> bool: + observations: Final = tuple(replica.observe() for replica in replicas) + assert not any(item.ready for item in observations), "Failed migration exposed a ready proxy" + return all(item.exit_code is not None for item in observations) + + until("every replica to reject startup", all_stopped) + for replica in replicas: + assert replica.state().ExitCode != 0, "Failed startup returned success" + assert marker in replica.logs(), f"Startup failed outside the expected migration: {marker}" + + +def waiting(replicas: tuple[Replica, ...], seconds: float) -> None: + deadline: Final = time.monotonic() + seconds + while time.monotonic() < deadline: + assert all(item.exit_code is None and not item.ready for item in (replica.observe() for replica in replicas)), ( + "Contending replica exited or served early" + ) + time.sleep(0.25) + + +@dataclass(frozen=True, slots=True) +class Containers: + image: str + output: Path + + @contextmanager + def start( + self, + database: Database, + migrations: tuple[Migration, ...] = (), + *, + v2: bool = True, + disabled: bool = False, + environment: Mapping[str, str] | None = None, + ) -> Generator[Replica]: + name: Final = f"litellm-migration-{uuid4().hex[:16]}" + directory: Final = self.output / name + directory.mkdir(parents=True) + for migration in migrations: + write_migration(directory, migration) + (directory / "config.yaml").write_text( + "model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n" + ) + env: Final = { + "DATABASE_URL": prisma_url(database.container_url, database.schema), + "LITELLM_MASTER_KEY": MASTER_KEY, + "LITELLM_SALT_KEY": MASTER_KEY, + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_TELEMETRY": "False", + "LITELLM_LOG": "INFO", + "DATABASE_CONNECTION_POOL_LIMIT": "2", + "DEFAULT_NUM_WORKERS_LITELLM_PROXY": "1", + "USE_V2_MIGRATION_RESOLVER": str(v2).lower(), + "DISABLE_SCHEMA_UPDATE": str(disabled).lower(), + "LITELLM_MIGRATION_DIR": "/migration-test/prisma", + "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT": "180", + **(environment or {}), + } + try: + docker( + "run", + "-d", + "--name", + name, + "--label", + "litellm-migration-test=true", + "--add-host", + "host.docker.internal:host-gateway", + "-p", + "127.0.0.1::4000", + "-v", + f"{directory}:/migration-test", + *(arg for key, value in env.items() for arg in ("-e", f"{key}={value}")), + self.image, + "--config", + "/migration-test/config.yaml", + "--host", + "0.0.0.0", + "--port", + "4000", + ) + port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1]) + replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory) + yield replica + finally: + try: + state: Final = subprocess.run( + ("docker", "inspect", "--format", "{{json .State}}", name), + capture_output=True, + text=True, + timeout=30, + ) + (directory / "state.json").write_text(state.stdout or state.stderr) + logs: Final = subprocess.run(("docker", "logs", name), capture_output=True, text=True, timeout=30) + (directory / "proxy.log").write_text(logs.stdout + logs.stderr) + finally: + subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True) + + +def write_migration(directory: Path, migration: Migration) -> None: + path: Final = directory / "prisma" / "migrations" / migration.name + path.mkdir(parents=True) + (path / "migration.sql").write_text(migration.script) diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py new file mode 100644 index 00000000000..a370c21ba0b --- /dev/null +++ b/tests/e2e/migrations/database.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Final, LiteralString +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit +from uuid import uuid4 + +import psycopg +from psycopg import sql +from pydantic import TypeAdapter + +Scalar = str | int | bool | None +ROWS: Final = TypeAdapter(tuple[tuple[Scalar, ...], ...]) +GATE_KEY: Final = 39178002 +PRISMA_LOCK: Final = 72707369 +COORDINATOR_LOCK: Final = int.from_bytes(b"llm_mig2", "big") + + +def connect_url(url: str, name: str) -> str: + return urlunsplit(urlsplit(url)._replace(path=f"/{name}", query="")) + + +def prisma_url(url: str, schema: str) -> str: + parsed: Final = urlsplit(url) + query: Final = tuple((key, value) for key, value in parse_qsl(parsed.query) if key != "schema") + return urlunsplit(parsed._replace(query=urlencode((*query, ("schema", schema))))) + + +@dataclass(frozen=True, slots=True) +class Database: + name: str + url: str + container_url: str + schema: str = "public" + + @contextmanager + def connection(self) -> Generator[psycopg.Connection[tuple[object, ...]]]: + with psycopg.connect(self.url, autocommit=True, connect_timeout=5) as connection: + connection.execute(sql.SQL("SET search_path TO {}").format(sql.Identifier(self.schema))) + connection.execute("SET statement_timeout = '15s'") + yield connection + + def execute(self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = ()) -> None: + with self.connection() as connection: + connection.execute(statement, params or None) + + def query( + self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = () + ) -> tuple[tuple[Scalar, ...], ...]: + with self.connection() as connection: + return ROWS.validate_python(connection.execute(statement, params or None).fetchall()) + + def exists(self, name: str) -> bool: + return self.query("SELECT to_regclass(%s) IS NOT NULL", (name,)) == ((True,),) + + def history(self) -> tuple[tuple[Scalar, ...], ...]: + if not self.exists("_prisma_migrations"): + return () + return self.query( + "SELECT id, migration_name, checksum, started_at::text, finished_at::text, rolled_back_at::text, " + "applied_steps_count, logs FROM _prisma_migrations ORDER BY id" + ) + + def blocked(self, key: int = GATE_KEY) -> tuple[tuple[Scalar, ...], ...]: + return self.query( + "SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND NOT granted " + "AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND classid = %s AND objid = %s ORDER BY pid", + (key >> 32, key & 0xFFFFFFFF), + ) + + @contextmanager + def lock(self, key: int = GATE_KEY) -> Generator[None]: + with self.connection() as connection: + connection.execute("SELECT pg_advisory_lock(%s)", (key,)) + try: + yield + finally: + connection.execute("SELECT pg_advisory_unlock(%s)", (key,)) + + +@dataclass(frozen=True, slots=True) +class Databases: + admin_url: str + container_admin_url: str + + @contextmanager + def create(self, template: Database | None = None, schema: str = "public") -> Generator[Database]: + name: Final = f"litellm_migration_test_{uuid4().hex[:20]}" + database: Final = Database( + name, connect_url(self.admin_url, name), connect_url(self.container_admin_url, name), schema + ) + with psycopg.connect(self.admin_url, autocommit=True, connect_timeout=5) as connection: + connection.execute( + sql.SQL("CREATE DATABASE {} TEMPLATE {}").format( + sql.Identifier(name), sql.Identifier(template.name if template else "template0") + ) + ) + try: + yield database + finally: + with psycopg.connect(self.admin_url, autocommit=True, connect_timeout=5) as connection: + connection.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) + + +@contextmanager +def restricted_user(database: Database) -> Generator[Database]: + role: Final = f"migration_reader_{uuid4().hex[:16]}" + password: Final = "migration-test-password" + with database.connection() as connection: + connection.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format(sql.Identifier(role), sql.Literal(password)) + ) + try: + database.execute( + sql.SQL("GRANT USAGE ON SCHEMA {} TO {}").format(sql.Identifier(database.schema), sql.Identifier(role)) + ) + database.execute( + sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA {} TO {}").format( + sql.Identifier(database.schema), sql.Identifier(role) + ) + ) + local: Final = urlsplit(database.url) + remote: Final = urlsplit(database.container_url) + yield Database( + database.name, + urlunsplit(local._replace(netloc=f"{role}:{password}@{local.hostname}:{local.port}")), + urlunsplit(remote._replace(netloc=f"{role}:{password}@{remote.hostname}:{remote.port}")), + database.schema, + ) + finally: + database.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role))) + database.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) diff --git a/tests/e2e/migrations/startup_models.py b/tests/e2e/migrations/startup_models.py new file mode 100644 index 00000000000..03a9b4eda78 --- /dev/null +++ b/tests/e2e/migrations/startup_models.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass + +from pydantic import BaseModel + + +class Readiness(BaseModel): + status: str = "" + db: str = "" + + +class ContainerState(BaseModel): + Running: bool + ExitCode: int + + +@dataclass(frozen=True, slots=True) +class Observation: + exit_code: int | None + ready: bool + + +@dataclass(frozen=True, slots=True) +class Migration: + name: str + script: str diff --git a/tests/e2e/migrations/test_legacy.py b/tests/e2e/migrations/test_legacy.py new file mode 100644 index 00000000000..7ba73eb82e0 --- /dev/null +++ b/tests/e2e/migrations/test_legacy.py @@ -0,0 +1,84 @@ +from contextlib import ExitStack +from dataclasses import replace +from typing import Final, Literal + +import pytest + +from .checks import COMPLETE, assert_completed, confirmed_history, assert_original_proof, start_replicas +from .containers import Containers, failed, ready +from .database import Database, Databases + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +def adopt_legacy(containers: Containers, database: Database) -> None: + count: Final = database.query("SELECT count(*) FROM _prisma_migrations")[0][0] + existing_keys: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token') + database.execute( + "INSERT INTO \"LiteLLM_ShadowEvalJob\" (id, group_id, target_id, router_name, judge_model, shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', 'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())" + ) + database.execute("DROP TABLE _prisma_migrations") + with ExitStack() as stack: + replicas: Final = start_replicas(stack, containers, database) + ready(replicas, database) + logs: Final = "\n".join(replica.logs() for replica in replicas) + for detail in ( + "Legacy migration history was missing", + "historical data backfills were not replayed or verified", + "Continuing startup", + ): + assert detail in logs + assert database.query("SELECT count(*) FROM _prisma_migrations") == ((count,),) + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL OR applied_steps_count <> 0" + ) == ((0,),) + assert set(existing_keys).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')) + assert database.query("SELECT stopped_by FROM \"LiteLLM_ShadowEvalJob\" WHERE id = 'migration-legacy'") == ( + (None,), + ) + + +class TestLegacyMigrations: + def test_matching_schema_warns_and_starts(self, containers: Containers, database: Database) -> None: + adopt_legacy(containers, database) + + @pytest.mark.parametrize("fault", ("schema_drift", "custom_migrations", "empty_ledger")) + def test_unrecognized_legacy_state_is_not_baselined( + self, containers: Containers, database: Database, fault: str + ) -> None: + if fault == "empty_ledger": + database.execute("TRUNCATE _prisma_migrations") + else: + database.execute("DROP TABLE _prisma_migrations") + if fault == "schema_drift": + database.execute('ALTER TABLE "LiteLLM_VerificationToken" DROP COLUMN key_alias CASCADE') + with containers.start(database, (COMPLETE,) if fault == "custom_migrations" else ()) as replica: + failed((replica,), "Cannot automatically baseline" if fault != "empty_ledger" else "migration") + assert not database.exists("migration_effect") + if database.exists("_prisma_migrations"): + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NOT NULL AND applied_steps_count <> 1" + ) == ((0,),) + + @pytest.mark.parametrize("scenario", ("upgrade", "recovery", "legacy")) + def test_non_default_schema( + self, containers: Containers, databases: Databases, scenario: Literal["upgrade", "recovery", "legacy"] + ) -> None: + with databases.create(schema="migration tenant") as database: + with containers.start(database) as seed: + ready((seed,), database) + match scenario: + case "upgrade": + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_completed(database) + case "recovery": + original: Final = confirmed_history(database) + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_original_proof(database, original, True) + case "legacy": + adopt_legacy(containers, database) + public: Final = replace(database, schema="public") + assert not public.exists("_prisma_migrations") + assert not public.exists('"LiteLLM_VerificationToken"') diff --git a/tests/e2e/migrations/test_pooling.py b/tests/e2e/migrations/test_pooling.py new file mode 100644 index 00000000000..e0a3693b33e --- /dev/null +++ b/tests/e2e/migrations/test_pooling.py @@ -0,0 +1,135 @@ +import subprocess +from collections.abc import Generator +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit, urlunsplit +from uuid import uuid4 + +import psycopg +import pytest +from psycopg import sql + +from .checks import COMPLETE, assert_completed +from .containers import Containers, docker, ready, until +from .database import Database, Databases, prisma_url, restricted_user + +POOL_IMAGE: Final = ( + "ghcr.io/cloudnative-pg/pgbouncer@sha256:e6ddfe22d845e603825e235dd8334b21ecd125abea2a2172478f556b8dee2bb8" +) +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +@contextmanager +def application_user(database: Database) -> Generator[Database]: + with restricted_user(database) as application: + role: Final = sql.Identifier(str(urlsplit(application.url).username)) + schema: Final = sql.Identifier(database.schema) + database.execute(sql.SQL("REVOKE CREATE ON SCHEMA {} FROM PUBLIC").format(schema)) + for statement in ( + "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {} TO {}", + "GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {} TO {}", + "ALTER DEFAULT PRIVILEGES IN SCHEMA {} GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO {}", + "ALTER DEFAULT PRIVILEGES IN SCHEMA {} GRANT USAGE, SELECT ON SEQUENCES TO {}", + ): + database.execute(sql.SQL(statement).format(schema, role)) + assert application.query("SELECT has_schema_privilege(current_user, %s, 'CREATE')", (database.schema,)) == ( + (False,), + ) + yield application + + +@contextmanager +def pool(database: Database, output: Path) -> Generator[str]: + name: Final = f"litellm-migration-pool-{uuid4().hex[:12]}" + url: Final = urlsplit(database.container_url) + directory: Final = output / name + directory.mkdir(parents=True) + (directory / "users.txt").write_text(f'"{url.username}" "{url.password}"\n') + (directory / "pgbouncer.ini").write_text( + f"[databases]\n* = host={url.hostname} port={url.port} user={url.username} password={url.password}\n" + "[pgbouncer]\nlisten_addr = 0.0.0.0\nlisten_port = 6432\nauth_type = trust\nauth_file = /pool/users.txt\n" + "pool_mode = transaction\ndefault_pool_size = 1\nreserve_pool_size = 0\nmax_client_conn = 100\n" + "max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = extra_float_digits,options\n" + ) + try: + docker( + "run", + "-d", + "--name", + name, + "--label", + "litellm-migration-test=true", + "--add-host", + "host.docker.internal:host-gateway", + "-p", + "0.0.0.0::6432", + "-v", + f"{directory}:/pool:ro", + "--entrypoint", + "/usr/bin/pgbouncer", + POOL_IMAGE, + "/pool/pgbouncer.ini", + ) + port: Final = int(docker("port", name, "6432/tcp").splitlines()[0].rsplit(":", 1)[1]) + local_url: Final = urlunsplit(url._replace(netloc=f"{url.username}:{url.password}@127.0.0.1:{port}")) + + def connected() -> bool: + try: + with psycopg.connect(local_url, autocommit=True, connect_timeout=2) as connection: + return connection.execute("SELECT 1").fetchone() == (1,) + except psycopg.Error: + return False + + until("PgBouncer ready", connected, 30) + yield local_url.replace("127.0.0.1", "host.docker.internal") + "?pgbouncer=true" + finally: + try: + logs: Final = subprocess.run(("docker", "logs", name), text=True, capture_output=True, timeout=30) + (directory / "pool.log").write_text(logs.stdout + logs.stderr) + finally: + subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True) + + +class TestMigrationPooling: + @pytest.mark.parametrize("scenario,replica_count", (("fresh", 3), ("upgrade", 3), ("legacy", 3), ("upgrade", 6))) + def test_direct_migrations_with_one_application_backend( + self, + containers: Containers, + databases: Databases, + migrated_template: Database, + scenario: str, + replica_count: int, + ) -> None: + with databases.create(None if scenario == "fresh" else migrated_template) as database: + if scenario == "legacy": + database.execute("DROP TABLE _prisma_migrations") + with ( + application_user(database) as application, + pool(application, containers.output) as pooled_url, + ExitStack() as stack, + ): + replicas: Final = tuple( + stack.enter_context( + containers.start( + database, + (COMPLETE,) if scenario == "upgrade" else (), + environment={ + "DATABASE_URL": prisma_url(pooled_url, database.schema), + "DIRECT_URL": database.container_url, + }, + ) + ) + for _ in range(replica_count) + ) + ready(replicas, database) + if scenario == "upgrade": + assert_completed(database) + if scenario == "legacy": + assert any( + "historical data backfills were not replayed or verified" in replica.logs() + for replica in replicas + ) + assert database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count <> 0") == ( + (0,), + ) diff --git a/tests/e2e/migrations/test_recovery.py b/tests/e2e/migrations/test_recovery.py new file mode 100644 index 00000000000..58bf5c348d6 --- /dev/null +++ b/tests/e2e/migrations/test_recovery.py @@ -0,0 +1,183 @@ +from contextlib import ExitStack +from typing import Final, Literal +from uuid import uuid4 + +import pytest + +from .checks import ( + COMPLETE, + FATAL, + GATED, + NEXT, + assert_completed, + confirmed_history, + interrupt_owner, + assert_original_proof, + pause_completion, + start_replicas, + unconfirmed, +) +from .containers import Containers, failed, ready, until, waiting +from .database import COORDINATOR_LOCK, GATE_KEY, Database +from .startup_models import Migration + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestMigrationRecovery: + @pytest.mark.parametrize("after_commit", (False, True)) + def test_container_owner_crash(self, containers: Containers, database: Database, after_commit: bool) -> None: + interrupt_owner(containers, database, after_commit, stop_database_session=False) + history: Final = database.history() + assert database.query( + "SELECT applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,) + ) == ((int(after_commit),),) + with ExitStack() as stack: + successors: Final = start_replicas(stack, containers, database, (COMPLETE if after_commit else GATED,)) + if after_commit: + ready(successors, database) + assert_completed(database) + else: + unconfirmed(successors, database) + assert database.history() == history + + @pytest.mark.parametrize("after_commit", (False, True)) + def test_owner_and_database_session_crash( + self, containers: Containers, database: Database, after_commit: bool + ) -> None: + interrupt_owner(containers, database, after_commit) + history: Final = database.history() + with ExitStack() as stack: + successors: Final = start_replicas(stack, containers, database, (COMPLETE,)) + if after_commit: + ready(successors, database) + assert_completed(database) + return + unconfirmed(successors, database) + assert database.history() == history + with containers.start(database, (COMPLETE,)) as restarted: + unconfirmed((restarted,), database) + assert database.history() == history + + @pytest.mark.parametrize("later_failure", (False, True)) + def test_remaining_migrations_after_recovery( + self, containers: Containers, database: Database, later_failure: bool + ) -> None: + original: Final = confirmed_history(database) + next_migration: Final = Migration( + NEXT.name, + f"SELECT pg_advisory_lock({GATE_KEY}); " + + (FATAL.script if later_failure else NEXT.script) + + f" SELECT pg_advisory_unlock({GATE_KEY});", + ) + with ExitStack() as stack: + with database.lock(): + owner: Final = stack.enter_context(containers.start(database, (COMPLETE, next_migration))) + + def pending() -> bool: + observation: Final = owner.observe() + assert observation.exit_code is None and not observation.ready, ( + "Recovered owner served before pending SQL completed" + ) + return bool(database.blocked()) + + until("recovering owner reached the next migration", pending) + assert_original_proof(database, original, True) + assert not database.exists("migration_next") + followers: Final = start_replicas(stack, containers, database, (COMPLETE, next_migration), count=2) + replicas: Final = (owner, *followers) + waiting(replicas, 1) + if later_failure: + failed(replicas, NEXT.name) + assert database.query( + "SELECT finished_at IS NULL, logs LIKE %s FROM _prisma_migrations WHERE migration_name = %s", + ("%MIGRATION_TEST_FATAL%", NEXT.name), + ) == ((True, True),) + else: + ready(replicas, database) + assert database.query("SELECT id FROM migration_next") == ((2,),) + assert_original_proof(database, original, True) + + def test_second_crash_during_recovery_is_atomic(self, containers: Containers, database: Database) -> None: + original: Final = confirmed_history(database) + pause_completion(database) + with database.lock(): + with containers.start(database, (COMPLETE,)) as recovering: + until("history update blocked before commit", lambda: bool(database.blocked())) + assert_original_proof(database, original, False) + blocked: Final = database.blocked() + assert len(blocked) == 1 + assert database.query("SELECT pg_terminate_backend(%s)", (blocked[0][0],)) == ((True,),) + failed((recovering,), "Lost or could not establish v2 migration coordination") + assert_original_proof(database, original, False) + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_original_proof(database, original, True) + + def test_competing_recovery_rechecks_stale_failures(self, containers: Containers, database: Database) -> None: + original: Final = confirmed_history(database) + with ExitStack() as stack: + with database.lock(COORDINATOR_LOCK): + replicas: Final = start_replicas(stack, containers, database, (COMPLETE,)) + until( + "all replicas observed the unfinished migration", + lambda: all( + "Waiting for the v2 migration coordinator lock" in replica.logs() for replica in replicas + ), + ) + assert_original_proof(database, original, False) + ready(replicas, database) + assert_original_proof(database, original, True) + + @pytest.mark.parametrize( + "fault", ("no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script") + ) + def test_unproven_history_is_never_repaired( + self, + containers: Containers, + database: Database, + fault: Literal["no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script"], + ) -> None: + confirmed_history(database) + match fault: + case "no_steps": + database.execute( + "UPDATE _prisma_migrations SET applied_steps_count = 0 WHERE migration_name = %s", (COMPLETE.name,) + ) + case "extra_steps": + database.execute( + "UPDATE _prisma_migrations SET applied_steps_count = 2 WHERE migration_name = %s", (COMPLETE.name,) + ) + case "failure_logs": + database.execute( + "UPDATE _prisma_migrations SET logs = 'permission denied' WHERE migration_name = %s", + (COMPLETE.name,), + ) + case "checksum": + database.execute( + "UPDATE _prisma_migrations SET checksum = %s WHERE migration_name = %s", ("0" * 64, COMPLETE.name) + ) + case "duplicate_history": + database.execute( + "INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) SELECT %s, migration_name, checksum, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + (str(uuid4()), COMPLETE.name), + ) + case "missing_script": + pass + history: Final = database.history() + with containers.start(database, () if fault == "missing_script" else (COMPLETE,)) as replica: + unconfirmed((replica,), database) + assert database.history() == history + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + def test_coordinator_timeout_preserves_proof(self, containers: Containers, database: Database) -> None: + original: Final = confirmed_history(database) + with database.lock(COORDINATOR_LOCK): + with containers.start( + database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "3"} + ) as replica: + failed((replica,), "Timed out waiting for another v2 migration resolver") + assert_original_proof(database, original, False) + with containers.start(database, (COMPLETE,)) as replica: + ready((replica,), database) + assert_original_proof(database, original, True) diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py new file mode 100644 index 00000000000..a648218cb26 --- /dev/null +++ b/tests/e2e/migrations/test_startup.py @@ -0,0 +1,100 @@ +from contextlib import ExitStack +from typing import Final + +import pytest + +from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas +from .containers import Containers, failed, ready, until, waiting +from .database import PRISMA_LOCK, Database, Databases, restricted_user +from .startup_models import Migration + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestMigrationStartup: + @pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False))) + def test_fresh_database(self, containers: Containers, databases: Databases, replicas: int, v2: bool) -> None: + with databases.create() as database, ExitStack() as stack: + ready(tuple(stack.enter_context(containers.start(database, v2=v2)) for _ in range(replicas)), database) + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL AND rolled_back_at IS NULL" + ) == ((0,),) + assert database.query("SELECT count(*) > 0 FROM _prisma_migrations") == ((True,),) + + def test_concurrent_upgrade(self, containers: Containers, database: Database) -> None: + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_completed(database) + + def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None: + with ExitStack() as stack: + with database.lock(): + owner: Final = stack.enter_context(containers.start(database, (GATED,))) + until("owner blocked in migration SQL", lambda: bool(database.blocked())) + followers: Final = start_replicas(stack, containers, database, (GATED,), count=2) + until("both followers attempted Prisma locking", lambda: len(database.blocked(PRISMA_LOCK)) == 2) + waiting((owner, *followers), 120) + ready((owner, *followers), database) + assert_completed(database, GATED) + + def test_lock_deadline_then_restart(self, containers: Containers, database: Database) -> None: + history: Final = database.history() + with database.lock(PRISMA_LOCK): + with containers.start( + database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"} + ) as replica: + until("Prisma lock contention", lambda: bool(database.blocked(PRISMA_LOCK))) + failed((replica,), "Timed out waiting for") + assert database.history() == history + assert not database.exists("migration_effect") + with containers.start(database, (COMPLETE,)) as restarted: + ready((restarted,), database) + assert_completed(database) + + def test_fatal_sql(self, containers: Containers, database: Database) -> None: + with ExitStack() as stack: + replicas: Final = start_replicas(stack, containers, database, (FATAL,)) + failed(replicas, COMPLETE.name) + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE %s AND finished_at IS NULL", + (COMPLETE.name, "%MIGRATION_TEST_FATAL%"), + ) == ((1,),) + + def test_duplicate_object_does_not_hide_incomplete_sql(self, containers: Containers, database: Database) -> None: + database.execute( + "CREATE TABLE migration_existing (id int PRIMARY KEY); INSERT INTO migration_existing VALUES (42)" + ) + migration: Final = Migration( + COMPLETE.name, "CREATE TABLE migration_existing (id int PRIMARY KEY); " + COMPLETE.script + ) + with ExitStack() as stack: + failed(start_replicas(stack, containers, database, (migration,)), COMPLETE.name) + assert not database.exists("migration_effect") + assert database.query("SELECT id FROM migration_existing") == ((42,),) + assert database.query( + "SELECT finished_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,) + ) == ((True,),) + + @pytest.mark.parametrize("v2", (True, False)) + def test_restart_preserves_history_and_data(self, containers: Containers, database: Database, v2: bool) -> None: + history: Final = database.history() + before: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token') + for _ in range(2): + with containers.start(database, v2=v2) as replica: + ready((replica,), database) + assert database.history() == history + assert set(before).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')) + + def test_disabled_migrations(self, containers: Containers, database: Database) -> None: + history: Final = database.history() + with containers.start(database, (FATAL,), disabled=True) as replica: + ready((replica,), database) + assert database.history() == history + + def test_insufficient_privileges(self, containers: Containers, database: Database) -> None: + history: Final = database.history() + with restricted_user(database) as limited: + with containers.start(limited, (COMPLETE,)) as replica: + failed((replica,), "permission denied") + assert database.history() == history + assert not database.exists("migration_effect") diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index e5826b18668..57133ea95c4 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -728,12 +728,36 @@ ERROR: relation "SomeTable" already exists """ +@pytest.mark.parametrize( + "pooled,direct,expected", + ( + ("postgresql://pool/db?pgbouncer=true", None, "postgresql://pool/db?pgbouncer=true"), + ("postgresql://pool/db?pgbouncer=true", "postgresql://writer/db", "postgresql://writer/db?schema=public"), + ( + "postgresql://pool/db?schema=tenant%20one&pgbouncer=true", + "postgresql://writer/db?sslmode=require&schema=wrong", + "postgresql://writer/db?sslmode=require&schema=tenant+one", + ), + ), +) +def test_v2_migrations_use_the_direct_connection_with_the_runtime_schema(pooled, direct, expected): + from litellm_proxy_extras.migration_lock import migration_environment + + environment = {"DATABASE_URL": pooled, "PRISMA_OFFLINE_MODE": "true"} + configured = {**environment, **({"DIRECT_URL": direct} if direct else {})} + migrated = migration_environment(configured) + + assert migrated["DATABASE_URL"] == expected + assert migrated["PRISMA_OFFLINE_MODE"] == "true" + assert configured["DATABASE_URL"] == pooled + + class _MigrateDeployHarness: """Drives _setup_database_v2 with a scripted sequence of `prisma migrate deploy` outcomes, with every recovery command faked out so nothing touches a database or the packaged migrations directory.""" - def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False): + def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False, confirmed_migrations=()): import subprocess as subprocess_module import litellm_proxy_extras.utils as utils_module @@ -744,16 +768,10 @@ class _MigrateDeployHarness: self._outcomes = list(outcomes) self._repeat_last = repeat_last self._subprocess_module = subprocess_module + self.confirmed_migrations = set(confirmed_migrations) monkeypatch.delenv("DATABASE_URL", raising=False) - monkeypatch.setattr( - ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path)) - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_create_baseline_migration", - staticmethod(self._fake_baseline), - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))) monkeypatch.setattr( ProxyExtrasDBManager, "_roll_back_migration", @@ -765,13 +783,15 @@ class _MigrateDeployHarness: staticmethod(self.resolved.append), ) monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) + monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {}) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) self.baseline_succeeds = True def _fake_baseline(self, *args, **kwargs): self.baselines += 1 - return self.baseline_succeeds + if not self.baseline_succeeds: + raise RuntimeError("The existing schema was not verified") def _next_outcome(self): if self._outcomes: @@ -791,79 +811,126 @@ class _MigrateDeployHarness: raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome) def run(self): - return ProxyExtrasDBManager._setup_database_v2(use_migrate=True) + while not ProxyExtrasDBManager._run_database_v2( + use_migrate=True, + recover_completed=self._fake_recovery, + baseline_existing=self._fake_baseline, + ): + continue + return True + + def _fake_recovery(self, name): + if name not in self.confirmed_migrations: + return False + self.confirmed_migrations.remove(name) + self.resolved.append(name) + return True class TestMigrateDeployAttemptAccounting: - """A `prisma db push` database has a full schema and no ledger, so the v2 - resolver baselines it and then works through every migration whose objects - already exist. Those recoveries make progress, so they must not spend the - retry budget, which is there to stop a run that is getting nowhere.""" - - def test_a_push_created_database_finishes_bootstrapping( - self, monkeypatch, tmp_path - ): - already_there = [ - "20250329084805_new_cron_job_table", - "20250806095134_rename_alias_to_server_name_mcp_table", - "20260224203854_add_agent_object_permissions_table", - "20260301120000_fourth_table", - "20260302120000_fifth_table", - "20260303120000_sixth_table", - ] + def test_a_push_created_database_finishes_bootstrapping(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness( monkeypatch, tmp_path, - [_P3005_STDERR] - + [_p3018_stderr(name) for name in already_there] - + ["ok"], + [_P3005_STDERR, "ok"], ) assert harness.run() is True assert harness.baselines == 1 - assert harness.resolved == already_there - assert len(harness.deploy_calls) == len(already_there) + 2 + assert harness.resolved == [] + assert len(harness.deploy_calls) == 2 - def test_repeated_recovery_of_one_migration_still_gives_up( - self, monkeypatch, tmp_path - ): + def test_repeated_recovery_of_one_migration_still_gives_up(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness( monkeypatch, tmp_path, [_p3018_stderr("20250329084805_new_cron_job_table")], repeat_last=True, + confirmed_migrations=("20250329084805_new_cron_job_table",), ) with pytest.raises(RuntimeError): harness.run() - assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1 + assert len(harness.deploy_calls) == 2 + assert harness.resolved == ["20250329084805_new_cron_job_table"] def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path): - harness = _MigrateDeployHarness( - monkeypatch, tmp_path, ["timeout"], repeat_last=True - ) + harness = _MigrateDeployHarness(monkeypatch, tmp_path, ["timeout"], repeat_last=True) with pytest.raises(RuntimeError): harness.run() assert len(harness.deploy_calls) == _ATTEMPT_BUDGET - def test_a_baseline_that_never_lands_stops_after_the_budget( - self, monkeypatch, tmp_path - ): - harness = _MigrateDeployHarness( - monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True - ) + def test_an_unverified_baseline_stops_without_replaying_migrations(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness(monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True) harness.baseline_succeeds = False with pytest.raises(RuntimeError): harness.run() - assert len(harness.deploy_calls) == _ATTEMPT_BUDGET + assert len(harness.deploy_calls) == 1 + + def test_lock_contention_does_not_spend_the_failure_budget(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + ["Error: P1002\nTimed out waiting for the advisory lock"] * 6 + ["ok"], + ) + assert harness.run() is True + assert len(harness.deploy_calls) == 7 + + def test_duplicate_object_error_without_completion_proof_is_fatal(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness(monkeypatch, tmp_path, [_p3018_stderr("20260101000000_x")]) + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + harness.run() + assert harness.resolved == [] + assert len(harness.deploy_calls) == 1 + + @pytest.mark.parametrize("name", ("20260101000000_x", "20260101000000_migration with spaces")) + def test_an_interrupted_migration_with_confirmed_sql_can_finish(self, monkeypatch, tmp_path, name): + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [f"Error: P3009\nThe `{name}` migration failed", "ok"], + confirmed_migrations=(name,), + ) + assert harness.run() is True + assert harness.resolved == [name] + + def test_an_interrupted_migration_without_confirmation_stops(self, monkeypatch, tmp_path): + name = "20260101000000_x" + started = "2026-09-12 20:15:06.694553 UTC" + report = f"Error: P3009\nThe `{name}` migration started at {started} failed" + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [report], + ) + with pytest.raises(RuntimeError, match="Migration completion could not be verified") as failure: + harness.run() + message = str(failure.value) + assert name in message + assert started in message + assert "start record but no successful completion record" in message + assert "cannot determine whether its SQL committed" in message + assert "avoid repeating or skipping database changes" in message + assert "_prisma_migrations" in message + assert "migration.sql" in message + assert "same database" in message + assert "Only after verifying every migration change is present" in message + assert "prisma migrate resolve --applied " in message + assert "Only after verifying no migration changes remain" in message + assert "prisma migrate resolve --rolled-back " in message + assert "leave migration history unchanged" in message + assert "Repeated restarts alone" in message + assert report in message + assert len(harness.deploy_calls) == 1 + assert harness.resolved == [] def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness( monkeypatch, tmp_path, - ["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"], + ['Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near "SLECT"\n'], repeat_last=True, ) @@ -873,6 +940,36 @@ class TestMigrateDeployAttemptAccounting: assert harness.resolved == [] +@pytest.mark.parametrize( + "steps,logs,script,expected", + ( + (1, "", b"CREATE TABLE item (id int);", True), + (0, "", b"CREATE TABLE item (id int);", False), + (0, "already exists", b"CREATE TABLE item (id int);", False), + (1, "permission denied", b"CREATE TABLE item (id int);", False), + (1, "", b"CREATE TABLE item (id text);", False), + (2, "", b"CREATE TABLE item (id int);", False), + ), +) +def test_migration_completion_requires_a_matching_successful_script(steps, logs, script, expected): + import hashlib + + from litellm_proxy_extras.migration_recovery import MigrationProgress + + progress = MigrationProgress(hashlib.sha256(b"CREATE TABLE item (id int);").hexdigest(), steps, logs) + assert progress.confirms_completion(script) is expected + + +def test_prisma_lock_waiting_has_its_own_deadline(): + from litellm_proxy_extras.utils import _MigrateAttemptBudget + + budget = _MigrateAttemptBudget(attempts_left=4, contention_seconds_left=2) + waiting = budget.after_contention(1) + assert waiting.attempts_left == 4 + with pytest.raises(RuntimeError, match="advisory lock"): + waiting.after_contention(2) + + class TestJWTKeyMappingCascade: """Regression tests for issue #33702. diff --git a/tests/proxy_migration_tests/test_migration_ci.py b/tests/proxy_migration_tests/test_migration_ci.py new file mode 100644 index 00000000000..30fe7383c07 --- /dev/null +++ b/tests/proxy_migration_tests/test_migration_ci.py @@ -0,0 +1,36 @@ +import importlib.util +from pathlib import Path +from typing import Final + +import pytest + +SCRIPT: Final = Path(__file__).resolve().parents[2] / ".circleci/scripts/run_migration_tests.py" +SPEC: Final = importlib.util.spec_from_file_location("migration_ci", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +MODULE: Final = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +@pytest.mark.parametrize( + "xml,expected,exit_code,passed", + ( + ('', 2, 0, True), + ('', 2, 0, False), + ("", 1, 0, False), + ("", 1, 0, False), + ("", 1, 0, False), + ("", 1, 1, False), + ("", 1, 5, False), + ("", 1, 0, False), + ("', 2, 0, False), + ), +) +def test_only_a_complete_passing_suite_can_certify_an_image( + tmp_path: Path, xml: str | None, expected: int, exit_code: int, passed: bool +) -> None: + path: Final = tmp_path / "results.xml" + if xml is not None: + path.write_text(xml) + assert MODULE.successful_junit(path, expected, exit_code) is passed From c44757fc010a0f81fe9c86fcabc43e95ecb8dd57 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:33:09 -0700 Subject: [PATCH 025/349] ci: fetch migration test revisions over HTTPS --- .circleci/config.yml | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index f6f31651306..c6e18c40213 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -22,11 +22,12 @@ commands: environment: MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >> command: | - if [ -n "$MIGRATION_SOURCE_SHA" ]; then - [[ "$MIGRATION_SOURCE_SHA" =~ ^[0-9a-f]{40}$ ]] || exit 1 - git fetch origin "$MIGRATION_SOURCE_SHA" - git checkout --detach "$MIGRATION_SOURCE_SHA" - fi + revision="${MIGRATION_SOURCE_SHA:-$CIRCLE_SHA1}" + [[ "$revision" =~ ^[0-9a-f]{40}$ ]] || exit 1 + git init + git remote add origin https://github.com/BerriAI/litellm.git + git fetch --depth 1 origin "$revision" + git checkout --detach FETCH_HEAD skip_if_unrelated_changes: parameters: category: @@ -2869,14 +2870,24 @@ jobs: destination: e2e-server-root-path-playwright-report build_docker_database_image: + parameters: + migration_qualification: + type: boolean + default: false machine: image: ubuntu-2204:2024.04.1 resource_class: large working_directory: ~/project steps: - - checkout - - checkout_migration_source - - skip_if_unrelated_changes + - when: + condition: << parameters.migration_qualification >> + steps: + - checkout_migration_source + - unless: + condition: << parameters.migration_qualification >> + steps: + - checkout + - skip_if_unrelated_changes - run: name: Build Docker image @@ -2923,7 +2934,6 @@ jobs: MIGRATION_TEST_OUTPUT: /tmp/migration-results PYTHONPATH: tests/e2e steps: - - checkout - checkout_migration_source - install_uv - install_rust @@ -3024,7 +3034,8 @@ workflows: migration_startup: when: << pipeline.parameters.run_migration_tests >> jobs: &migration_jobs - - build_docker_database_image + - build_docker_database_image: + migration_qualification: true - migration_startup_tests: name: migration-startup suite: startup From a37f0b4f544513d48001f1b2a86bd1eef59ca2c9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:49:58 -0700 Subject: [PATCH 026/349] test: isolate migration CI selection and exercise resolver boundaries --- .github/e2e-stack/select_tests.py | 2 +- .../litellm_proxy_extras/utils.py | 9 +- .../tests/test_setup_database_fail_fast.py | 156 +++++++----------- .../test_e2e_changed_gate.py | 5 + tests/e2e/conftest.py | 4 +- tests/e2e/migrations/checks.py | 12 +- tests/e2e/migrations/test_legacy.py | 7 +- tests/e2e/migrations/test_pooling.py | 3 +- tests/e2e/migrations/test_recovery.py | 4 +- tests/e2e/migrations/test_startup.py | 3 +- .../test_litellm_proxy_extras_utils.py | 12 +- 11 files changed, 93 insertions(+), 124 deletions(-) diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 238818a0d36..9dd880c05cd 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -4,7 +4,7 @@ from typing import Final SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$") UNSUPPORTED: Final = re.compile( - r"^tests/e2e/(ui|claude_code|load)/" + r"^tests/e2e/(ui|claude_code|load|migrations)/" r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$" r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$" diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2749db5d754..8a83c786e02 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -859,7 +859,11 @@ class ProxyExtrasDBManager: def baseline_existing(migrations_dir: str) -> None: with migration_lock(lock_url) as coordinator: baseline_current_schema( - coordinator, schema, Path(migrations_dir), _get_prisma_command(), migration_environment(_get_prisma_env()) + coordinator, + schema, + Path(migrations_dir), + _get_prisma_command(), + migration_environment(_get_prisma_env()), ) while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing): @@ -1054,7 +1058,8 @@ class ProxyExtrasDBManager: migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr: logger.info( - "Migration %s deadlocked against a concurrent migrate deploy, rolling its ledger row back and retrying", + "Migration %s deadlocked against a concurrent migrate deploy, " + "rolling its ledger row back and retrying", migration_name, ) ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 338c571eb4f..832075f6fbe 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -6,7 +6,8 @@ The v2 resolver is opt-in via `--use_v2_migration_resolver` / the """ import subprocess -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch import pytest @@ -31,10 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str): def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): """v2: a permission failure during migrate deploy raises RuntimeError.""" - monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + _stub_v2_env(monkeypatch, tmp_path) stderr = ( "Error: P3018\nMigration name: 20250326162113_baseline\n" @@ -47,10 +45,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): """v2: a non-idempotent migration failure raises (no silent recovery).""" - monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + _stub_v2_env(monkeypatch, tmp_path) stderr = ( "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n" @@ -131,8 +126,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path): def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path): """v2: a failing `prisma db push` must raise RuntimeError, not leak CalledProcessError past proxy_cli.py's `except RuntimeError`.""" - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) stderr = "db push error" with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): @@ -149,8 +143,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): import psycopg monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) class _FakeConn: def __enter__(self): @@ -173,18 +166,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path): - _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "relation already exists") - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: pytest.fail("duplicate-object errors do not prove rollback is safe"), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("duplicate-object errors do not prove all SQL completed"), - ) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs="relation already exists") stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists" with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) @@ -197,28 +179,15 @@ def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path): def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): - """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + _stub_v2_env(monkeypatch, tmp_path) + run = Mock(side_effect=_succeed_after(0, "")) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run) - class FakeResult: - stdout = "Applied migration.\n" - stderr = "" - - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult()) - - resolve_called = {"n": 0} - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_all_migrations", - lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1), + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) is True + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], ) - ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) - assert ok is True - assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" - _DEADLOCK_P3018_STDERR = ( "Error: P3018\n" @@ -228,12 +197,34 @@ _DEADLOCK_P3018_STDERR = ( ) -def _stub_v2_env(monkeypatch, tmp_path): +def _stub_v2_env(monkeypatch, tmp_path, ledger_logs=""): + import psycopg + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) monkeypatch.setattr("time.sleep", lambda _: None) + connection = MagicMock() + connection.__enter__.return_value = connection + cursor = connection.cursor.return_value.__enter__.return_value + cursor.execute.return_value = cursor + cursor.fetchone.return_value = SimpleNamespace(acquired=True) + cursor.fetchall.return_value = [] + empty = MagicMock() + empty.fetchall.return_value = [] + empty.fetchone.return_value = None + ledger = MagicMock() + ledger.fetchone.return_value = (ledger_logs,) + + def execute(query, *args, **kwargs): + if "SELECT logs FROM" in str(query): + if ledger_logs is None: + raise psycopg.OperationalError("ledger is unavailable") + return ledger + return empty + + connection.execute.side_effect = execute + monkeypatch.setattr("psycopg.connect", lambda *args, **kwargs: connection) def _succeed_after(failures: int, stderr: str): @@ -259,28 +250,21 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): instance rolls the ledger row back and retries instead of dying.""" _stub_v2_env(monkeypatch, tmp_path) - rolled_back = [] - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: rolled_back.append(name), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("a deadlocked migration must never be marked applied"), - ) - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR)) + run = Mock(side_effect=_succeed_after(1, _DEADLOCK_P3018_STDERR)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True - assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], + ["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"], + ["migrate", "deploy"], + ) def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): """v2: a deadlock on every attempt still fails after the retry budget.""" _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr(ProxyExtrasDBManager, "_v2_roll_back_migration_best_effort", lambda name: None) with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", @@ -293,7 +277,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path): """v2: the surviving instance sees the victim's failed ledger row as P3009. When that row's logs show a deadlock, roll it back and retry.""" - _stub_v2_env(monkeypatch, tmp_path) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs="ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock") stderr = ( "Error: P3009\n" @@ -301,27 +285,16 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ "The `20260415120000_health_check_latest_per_model_index` migration " "started at 2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_failed_migration_logs", - lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock", - ) - rolled_back = [] - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: rolled_back.append(name), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("a deadlocked migration must never be marked applied"), - ) - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) + run = Mock(side_effect=_succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True - assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], + ["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"], + ["migrate", "deploy"], + ) def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path): @@ -332,12 +305,6 @@ def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_pat "The `20260415120000_health_check_latest_per_model_index` migration " "started at 2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "") - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: pytest.fail("empty logs do not prove rollback is safe"), - ) with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) ) as run: @@ -350,7 +317,7 @@ def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_pat def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): """v2: an unreadable ledger cannot establish that P3009 was a deadlock.""" - _stub_v2_env(monkeypatch, tmp_path) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs=None) stderr = ( "Error: P3009\n" @@ -358,12 +325,6 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): "The `20260415120000_health_check_latest_per_model_index` migration " "started at 2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), - ) monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) with pytest.raises(RuntimeError, match="Migration completion could not be verified"): @@ -372,7 +333,7 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): """v2: a failed ledger row whose logs show a real SQL error stays fatal.""" - _stub_v2_env(monkeypatch, tmp_path) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs='ERROR: syntax error at or near "BRKN"') stderr = ( "Error: P3009\n" @@ -380,11 +341,6 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): "The `20260101000000_genuinely_broken` migration started at " "2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_failed_migration_logs", - lambda name: 'ERROR: syntax error at or near "BRKN"', - ) with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="Migration completion could not be verified"): diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 588402e3996..a628fbb0633 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -112,6 +112,7 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: ( (("tests/e2e/logging/test_datadog_e2e.py", "litellm/router.py"), ("tests/e2e/logging/test_datadog_e2e.py",)), (("tests/e2e/ui/test_keys.py", "tests/e2e/claude_code/test_cli.py", "tests/e2e/load/test_burst.py"), ()), + (("tests/e2e/migrations/test_startup.py", "tests/e2e/migrations/test_recovery.py"), ()), (("tests/e2e/batches/test_managed_files_enforcement_e2e.py",), ()), (("tests/e2e/guardrails/test_presidio_masking_e2e.py",), ()), (("tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py",), ()), @@ -157,6 +158,10 @@ def test_a_changed_canary_file_is_selected_once_alongside_a_harness_change() -> assert select_tests((CANARY[1], "tests/e2e/proxy_client.py")) == CANARY +def test_dedicated_migration_tests_do_not_suppress_shared_harness_canaries() -> None: + assert select_tests(("tests/e2e/migrations/test_startup.py", "tests/e2e/conftest.py")) == CANARY + + def test_the_canary_joins_directly_selected_files_in_sorted_order() -> None: assert select_tests(("tests/e2e/logging/test_datadog_e2e.py", ".github/e2e-stack/up.sh")) == ( *CANARY, diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4a5f0aa880f..53d35effdc9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -64,7 +64,9 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line("markers", "migration_startup: isolated container startup tests run by the migration CI workflow") + config.addinivalue_line( + "markers", "migration_startup: isolated container startup tests run by the migration CI workflow" + ) config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py index b14631ccad8..619ad3b0e6c 100644 --- a/tests/e2e/migrations/checks.py +++ b/tests/e2e/migrations/checks.py @@ -29,7 +29,8 @@ def start_replicas( def assert_completed(database: Database, migration: Migration = COMPLETE) -> None: assert database.query( - "SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + 'SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM ' + '_prisma_migrations WHERE migration_name = %s', (migration.name,), ) == ((True, True, 1),), "Expected exactly one successful SQL execution" assert database.query("SELECT id FROM migration_effect") == ((1,),) @@ -47,7 +48,8 @@ def confirmed_history(database: Database) -> str: def assert_original_proof(database: Database, row_id: str, finished: bool) -> None: assert database.query( - "SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", + 'SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM ' + '_prisma_migrations WHERE migration_name = %s', (COMPLETE.name,), ) == ((row_id, 1, finished, True),), "Recovery lost or replaced the original durable SQL proof" assert database.query("SELECT id FROM migration_effect") == ((1,),) @@ -59,7 +61,8 @@ def pause_completion(database: Database) -> None: "CREATE FUNCTION migration_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " "IF NEW.migration_name = {name} AND NEW.finished_at IS NOT NULL THEN " "PERFORM pg_advisory_lock({gate}); PERFORM pg_advisory_unlock({gate}); END IF; RETURN NEW; END $$; " - "CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW EXECUTE FUNCTION migration_pause()" + 'CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW ' + 'EXECUTE FUNCTION migration_pause()' ).format(name=sql.Literal(COMPLETE.name), gate=sql.Literal(GATE_KEY)) ) @@ -105,7 +108,8 @@ def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None: failed(replicas, "Migration completion could not be verified") started: Final = str( database.query( - "SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM _prisma_migrations WHERE migration_name = %s", + "SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM " + '_prisma_migrations WHERE migration_name = %s', (COMPLETE.name,), )[0][0] ) diff --git a/tests/e2e/migrations/test_legacy.py b/tests/e2e/migrations/test_legacy.py index 7ba73eb82e0..ba5e77a3070 100644 --- a/tests/e2e/migrations/test_legacy.py +++ b/tests/e2e/migrations/test_legacy.py @@ -15,7 +15,9 @@ def adopt_legacy(containers: Containers, database: Database) -> None: count: Final = database.query("SELECT count(*) FROM _prisma_migrations")[0][0] existing_keys: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token') database.execute( - "INSERT INTO \"LiteLLM_ShadowEvalJob\" (id, group_id, target_id, router_name, judge_model, shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', 'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())" + 'INSERT INTO "LiteLLM_ShadowEvalJob" (id, group_id, target_id, router_name, judge_model, ' + "shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', " + "'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())" ) database.execute("DROP TABLE _prisma_migrations") with ExitStack() as stack: @@ -30,7 +32,8 @@ def adopt_legacy(containers: Containers, database: Database) -> None: assert detail in logs assert database.query("SELECT count(*) FROM _prisma_migrations") == ((count,),) assert database.query( - "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL OR applied_steps_count <> 0" + 'SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS ' + 'NOT NULL OR applied_steps_count <> 0' ) == ((0,),) assert set(existing_keys).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')) assert database.query("SELECT stopped_by FROM \"LiteLLM_ShadowEvalJob\" WHERE id = 'migration-legacy'") == ( diff --git a/tests/e2e/migrations/test_pooling.py b/tests/e2e/migrations/test_pooling.py index e0a3693b33e..c4015549de3 100644 --- a/tests/e2e/migrations/test_pooling.py +++ b/tests/e2e/migrations/test_pooling.py @@ -50,7 +50,8 @@ def pool(database: Database, output: Path) -> Generator[str]: f"[databases]\n* = host={url.hostname} port={url.port} user={url.username} password={url.password}\n" "[pgbouncer]\nlisten_addr = 0.0.0.0\nlisten_port = 6432\nauth_type = trust\nauth_file = /pool/users.txt\n" "pool_mode = transaction\ndefault_pool_size = 1\nreserve_pool_size = 0\nmax_client_conn = 100\n" - "max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = extra_float_digits,options\n" + 'max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = ' + 'extra_float_digits,options\n' ) try: docker( diff --git a/tests/e2e/migrations/test_recovery.py b/tests/e2e/migrations/test_recovery.py index 58bf5c348d6..80e5747eaac 100644 --- a/tests/e2e/migrations/test_recovery.py +++ b/tests/e2e/migrations/test_recovery.py @@ -159,7 +159,9 @@ class TestMigrationRecovery: ) case "duplicate_history": database.execute( - "INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) SELECT %s, migration_name, checksum, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + 'INSERT INTO _prisma_migrations (id, migration_name, checksum, ' + 'applied_steps_count) SELECT %s, migration_name, checksum, ' + 'applied_steps_count FROM _prisma_migrations WHERE migration_name = %s', (str(uuid4()), COMPLETE.name), ) case "missing_script": diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py index a648218cb26..dc628a8ed7a 100644 --- a/tests/e2e/migrations/test_startup.py +++ b/tests/e2e/migrations/test_startup.py @@ -56,7 +56,8 @@ class TestMigrationStartup: replicas: Final = start_replicas(stack, containers, database, (FATAL,)) failed(replicas, COMPLETE.name) assert database.query( - "SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE %s AND finished_at IS NULL", + 'SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE ' + '%s AND finished_at IS NULL', (COMPLETE.name, "%MIGRATION_TEST_FATAL%"), ) == ((1,),) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 57133ea95c4..bb329264a11 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -771,17 +771,7 @@ class _MigrateDeployHarness: self.confirmed_migrations = set(confirmed_migrations) monkeypatch.delenv("DATABASE_URL", raising=False) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_roll_back_migration", - staticmethod(lambda name: None), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - staticmethod(self.resolved.append), - ) + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {}) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) From 20d80b5420508c73391cca91be232b7f74041d1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:45:07 -0700 Subject: [PATCH 027/349] fix(guardrails): scan each choice's tool-call arguments apart on n>1 streams and log why a rewrite was discarded The rebuilt streamed response keyed tool-call fragments by tool index alone, so on n>1 chat streams the two choices' argument fragments were concatenated into one string and post_call guardrails scanned garbled JSON. Fragments are now keyed by (choice index, tool index). When a guardrail's rewrite cannot be written back to the stream (multi-choice streams, a rewrite that adds or drops a tool call, legacy-hook shapes the translation cannot rescan), the pipeline now logs a warning naming the guardrail and the exact reason before releasing the original stream. Also commits the regenerated dashboard API types that make check produced. --- .../streaming_chunk_builder_utils.py | 38 +++++---- .../chat/guardrail_translation/handler.py | 11 ++- .../chat/guardrail_translation/handler.py | 26 +++++-- .../guardrail_translation/handler.py | 11 ++- .../proxy/policy_engine/pipeline_executor.py | 78 ++++++++++++++----- .../test_streaming_chunk_builder_utils.py | 55 +++++++++++++ .../test_openai_guardrail_handler.py | 57 +++++++++++++- .../policy_engine/test_pipeline_executor.py | 42 ++++++---- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 9 files changed, 254 insertions(+), 66 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 90698296142..5ffe36573d5 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -138,9 +138,13 @@ class _ToolCallDelta(TypedDict, total=False): class _ToolCallChoice(TypedDict, total=False): + index: ReadOnly[int] delta: ReadOnly[_ToolCallDelta] +_ToolCallKey: TypeAlias = tuple[int, int] + + class _ToolCallChunk(TypedDict): choices: ReadOnly[Sequence[_ToolCallChoice]] @@ -416,40 +420,41 @@ class ChunkProcessor: @staticmethod def _iter_tool_call_fragments( tool_call_chunks: Sequence["_ToolCallChunk"], - ) -> Iterator[tuple[int, str, str]]: + ) -> Iterator[tuple[_ToolCallKey, str, str]]: for chunk in tool_call_chunks: for choice in chunk["choices"]: delta = choice.get("delta") if not delta: continue + choice_index = choice.get("index", 0) for tool_call in delta.get("tool_calls", ()): if not tool_call: continue if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + key = (choice_index, tool_call.get("index", 0)) function = tool_call.get("function") if isinstance(function, dict): if fragment_arguments := function.get("arguments"): - yield index, "arguments", fragment_arguments + yield key, "arguments", fragment_arguments elif function_arguments := getattr(function, "arguments", None): - yield index, "arguments", function_arguments + yield key, "arguments", function_arguments custom = tool_call.get("custom") if isinstance(custom, dict) and (custom_input := custom.get("input")): - yield index, "custom_input", custom_input + yield key, "custom_input", custom_input else: - index = getattr(tool_call, "index", 0) + key = (choice_index, getattr(tool_call, "index", 0)) function = getattr(tool_call, "function", None) if object_arguments := getattr(function, "arguments", None): - yield index, "arguments", object_arguments + yield key, "arguments", object_arguments custom = getattr(tool_call, "custom", None) if object_custom_input := getattr(custom, "input", None): - yield index, "custom_input", object_custom_input + yield key, "custom_input", object_custom_input @staticmethod - def _join_fragments_by_index_and_field( - fragment_records: Iterator[tuple[int, str, str]], - ) -> Mapping[tuple[int, str], str]: - def group_key(record: tuple[int, str, str]) -> tuple[int, str]: + def _join_fragments_by_key_and_field( + fragment_records: Iterator[tuple[_ToolCallKey, str, str]], + ) -> Mapping[tuple[_ToolCallKey, str], str]: + def group_key(record: tuple[_ToolCallKey, str, str]) -> tuple[_ToolCallKey, str]: return record[0], record[1] return MappingProxyType( @@ -467,13 +472,14 @@ class ChunkProcessor: tool_calls_list: list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ] = [] # mutable-ok: see return type - tool_call_map: Final[dict[int, dict[str, Any]]] = {} # Map to store tool calls by index + tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {} # Map to store tool calls by choice and index for chunk in tool_call_chunks: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) tool_calls = delta.get("tool_calls", []) + choice_index = choice.get("index", 0) for tool_call in tool_calls: # Handle both dict and object formats @@ -495,9 +501,9 @@ class ChunkProcessor: # Get index (handle both dict and object) if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + index = (choice_index, tool_call.get("index", 0)) else: - index = getattr(tool_call, "index", 0) + index = (choice_index, getattr(tool_call, "index", 0)) if index not in tool_call_map: tool_call_map[index] = { @@ -572,7 +578,7 @@ class ChunkProcessor: if isinstance(provider_fields, dict): merged_provider_fields.update(provider_fields) - joined_fragments: Final = self._join_fragments_by_index_and_field( + joined_fragments: Final = self._join_fragments_by_key_and_field( self._iter_tool_call_fragments(tool_call_chunks) ) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..c7ba5daec56 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1172,7 +1172,10 @@ class AnthropicMessagesHandler(BaseTranslation): if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + "the stream never reported a stop_reason, so the text rewrite has no assembled response to land on", + ) return responses_so_far def _prepare_request_data( @@ -1318,7 +1321,11 @@ class AnthropicMessagesHandler(BaseTranslation): if len(block_indices) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " + f"{len(block_indices)} tool_use blocks", + ) rewrites_by_block: Final = MappingProxyType( { index: after diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 58ff03e6a0d..e4943690639 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1041,13 +1041,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice.index for response in responses_so_far for choice in response.choices ) if len(stream_choice_indices) != 1: - # stream_chunk_builder collapses every choice into one index-0 - # choice, so a rewrite of the rebuilt response cannot be attributed - # back to a single choice on an n>1 stream: report it undeliverable - # rather than deliver the rewrite on the wrong choice from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and the rebuilt response's text rewrite " + "cannot be attributed to one of them", + ) target_choice_index: Final = next(iter(stream_choice_indices)) await self._apply_guardrail_responses_to_output_streaming( responses=responses_so_far, @@ -1105,10 +1105,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice.index for response in responses_so_far for choice in response.choices ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) - if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + if len(stream_choice_indices) != 1: from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " + "back on single-choice streams", + ) + if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " + f"{len(fragments_by_tool_call)}", + ) for before, (name, arguments), fragments in zip( pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call ): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..2be36a826f7 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -957,7 +957,10 @@ class OpenAIResponsesHandler(BaseTranslation): if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + "the stream carried no terminal response envelope to write the text rewrite back into", + ) return responses_so_far @staticmethod @@ -1070,7 +1073,11 @@ class OpenAIResponsesHandler(BaseTranslation): ): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls and the stream's " + f"{len(tool_call_items)} function_call items could not be lined up with them by call_id", + ) for output_item, rewrite in ( (output_item, rewrites_by_call_id[call_id]) for output_item, call_id in zip(tool_call_items, call_ids) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ad45781d5d2..26dd806b3e5 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -50,12 +50,13 @@ except ImportError: class UndeliverableStreamRewrite(Exception): - def __init__(self, guardrail_name: str) -> None: + def __init__(self, guardrail_name: str, reason: str) -> None: super().__init__( - f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's " - "streaming pipeline cannot deliver" + f"Guardrail '{guardrail_name}' rewrote the streamed response but the rewrite cannot be written " + f"back to the stream: {reason}" ) self.guardrail_name: Final = guardrail_name + self.reason: Final = reason class UnappliableRequestRewrite(Exception): @@ -91,8 +92,22 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent -def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: - return sent is not None and returned is not None and len(returned) != len(sent) +def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> tuple[int, int] | None: + if sent is None or returned is None or len(returned) == len(sent): + return None + return (len(sent), len(returned)) + + +def _tool_call_mismatch_reason( + sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None +) -> str | None: + if sent == returned: + return None + sent_count: Final = len(sent or ()) + returned_count: Final = len(returned or ()) + if sent_count == returned_count: + return "the legacy hook changed a tool call's name or arguments, which this path cannot write back" + return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}" _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) @@ -119,7 +134,7 @@ class _StreamRewriteObserver(CustomGuardrail): self.inner: Final = inner self.rewrote_texts = False self.rewrote_tool_calls = False - self.changed_tool_call_count = False + self.tool_call_count_change: tuple[int, int] | None = None def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -140,11 +155,22 @@ class _StreamRewriteObserver(CustomGuardrail): returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) - self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + self.tool_call_count_change = self.tool_call_count_change or _count_change( sent_tool_shapes, returned_tool_shapes ) return outputs + def discard_reason(self, deliver_rewrites: bool) -> str | None: + if self.tool_call_count_change is not None: + sent, returned = self.tool_call_count_change + return ( + f"the guardrail returned {returned} tool calls for a stream that carried {sent}, and a rewrite " + "that drops or adds a tool call cannot be written back" + ) + if not deliver_rewrites and (self.rewrote_texts or self.rewrote_tool_calls): + return "this endpoint's streaming pipeline does not write ended-stream rewrites back yet" + return None + class _ScannedTextRecorder(CustomGuardrail): def __init__(self, guardrail_name: str) -> None: @@ -209,13 +235,24 @@ class _LegacyHookStreamAdapter(CustomGuardrail): if rewrite is None: return inputs rescanned: Final = await self._rescan(rewrite, logging_obj) + guardrail_name: Final = self.guardrail_name or "unknown" if rescanned is None: - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_name, "the legacy hook's response could not be rescanned by this endpoint's translation" + ) rewritten: Final = rescanned.get("texts") - if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") - if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + returned_text_count: Final = len(_scanned_texts(rewritten)) + sent_text_count: Final = len(_scanned_texts(inputs.get("texts"))) + if returned_text_count != sent_text_count: + raise UndeliverableStreamRewrite( + guardrail_name, + f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}", + ) + tool_call_mismatch: Final = _tool_call_mismatch_reason( + _tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls")) + ) + if tool_call_mismatch is not None: + raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch) if not rewritten: return inputs rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} @@ -262,14 +299,16 @@ def _prepare_hook_input( def _release_original_chunks( guardrail_name: str, + reason: str, streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place originals: Sequence[object], ) -> None: streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives verbose_proxy_logger.warning( - "Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming " - "pipeline cannot deliver yet; the rewrite was discarded and the original stream released", + "Pipeline: guardrail '%s' rewrote the streamed response but the rewrite could not be written back to " + "the stream: %s. The whole rewrite, text rewrites included, was discarded and the original stream released", guardrail_name, + reason, ) @@ -442,13 +481,12 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, request_data=hook_input, ) - except UndeliverableStreamRewrite: - _release_original_chunks(step.guardrail, streaming_chunks, originals) + except UndeliverableStreamRewrite as undeliverable: + _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return - if observer.changed_tool_call_count or ( - not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) - ): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + discard_reason: Final = observer.discard_reason(deliver_rewrites) + if discard_reason is not None: + _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals) return if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index efe4209c1c9..2266258bf20 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1288,6 +1288,61 @@ def _tool_call_delta_chunk(tool_call: dict[str, object] | ChatCompletionDeltaToo return {"choices": [{"delta": {"tool_calls": [tool_call]}}]} +def _choice_tool_call_delta_chunk(choice_index: int, tool_call: dict[str, object]) -> dict[str, object]: + return {"choices": [{"index": choice_index, "delta": {"tool_calls": [tool_call]}}]} + + +def test_get_combined_tool_content_keeps_each_choices_arguments_apart_when_choices_share_a_tool_index(): + processor = ChunkProcessor.__new__(ChunkProcessor) + chunks = [ + _choice_tool_call_delta_chunk(0, {"index": 0, "id": "call_a", "type": "function", "function": {"name": "f"}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "id": "call_b", "type": "function", "function": {"name": "f"}}), + _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": '{"fruit": "pers'}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": '{"fruit": "dur'}}), + _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": 'immon"}'}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": 'ian"}'}}), + ] + + combined = processor.get_combined_tool_content(chunks) + + assert [(tool_call.id, tool_call.function.arguments) for tool_call in combined] == [ + ("call_a", '{"fruit": "persimmon"}'), + ("call_b", '{"fruit": "durian"}'), + ] + + +def test_stream_chunk_builder_keeps_each_choices_tool_call_arguments_apart(): + def chunk(choice_index: int, tool_call: ChatCompletionDeltaToolCall) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + object="chat.completion.chunk", + created=1234567890, + model="gpt-4.1-mini", + choices=[StreamingChoices(index=choice_index, delta=Delta(tool_calls=[tool_call]), finish_reason=None)], + ) + + def fragment(arguments: str, name: str | None = None, call_id: str | None = None) -> ChatCompletionDeltaToolCall: + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + response = stream_chunk_builder( + chunks=[ + chunk(0, fragment("", name="lookup_fruit", call_id="call_a")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_b")), + chunk(0, fragment('{"fruit": "pers')), + chunk(1, fragment('{"fruit": "dur')), + chunk(0, fragment('immon"}')), + chunk(1, fragment('ian"}')), + ] + ) + + assert [(tool_call.id, tool_call.function.arguments) for tool_call in response.choices[0].message.tool_calls] == [ + ("call_a", '{"fruit": "persimmon"}'), + ("call_b", '{"fruit": "durian"}'), + ] + + def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_order(): processor = ChunkProcessor.__new__(ChunkProcessor) first_fragments = [f"a{i};" for i in range(300)] diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5a29a96829f..60a5752e83a 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1267,7 +1267,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: handler = OpenAIChatCompletionsHandler() chunks = self._two_choice_stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: await handler.process_output_streaming_response( responses_so_far=chunks, guardrail_to_apply=self._world_masking_guardrail(), @@ -1275,6 +1275,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + assert raised.value.guardrail_name == "test-mask" + assert raised.value.reason == ( + "the stream carries 2 choices and the rebuilt response's text rewrite cannot be attributed to one of them" + ) + @staticmethod def _two_choice_tool_call_stream_chunks() -> list: from litellm.types.utils import ( @@ -1310,12 +1315,51 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: return [ chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), - chunk(0, fragment('{"fruit": "persimmon"}')), - chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, fragment('{"fruit": "pers')), + chunk(1, fragment('{"fruit": "dur')), + chunk(0, fragment('immon"}')), + chunk(1, fragment('ian"}')), chunk(0, None, finish_reason="tool_calls"), chunk(1, None, finish_reason="tool_calls"), ] + @staticmethod + def _recording_guardrail() -> CustomGuardrail: + class Recorder(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="recorder") + self.seen_inputs: list[GenericGuardrailAPIInputs] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.seen_inputs.append(inputs) + return inputs + + return Recorder() + + @pytest.mark.asyncio + async def test_ended_multi_choice_stream_scans_each_choices_tool_call_arguments_apart(self): + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + guardrail = self._recording_guardrail() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert [ + (tool_call["id"], tool_call["function"]["arguments"]) + for tool_call in guardrail.seen_inputs[-1]["tool_calls"] + ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1323,7 +1367,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: handler = OpenAIChatCompletionsHandler() chunks = self._two_choice_tool_call_stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: await handler.process_output_streaming_response( responses_so_far=chunks, guardrail_to_apply=MockGuardrail(guardrail_name="test"), @@ -1331,6 +1375,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + assert raised.value.guardrail_name == "test" + assert raised.value.reason == ( + "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" + ) + @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 0a2641082dc..e7689cc7d0c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1122,7 +1122,7 @@ class _RefusingTranslation: deliver_ended_stream_rewrites=False, ): responses_so_far[0]["text"] = "half-written" - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name) + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name, "the translation refused it") def _chunk(): @@ -1143,10 +1143,22 @@ async def _run_streaming_step(translation, streaming_chunks=None): ) -def _assert_passed_with_discard_warning(result, caplog): +NO_WRITE_BACK_REASON = "this endpoint's streaming pipeline does not write ended-stream rewrites back yet" + + +def _assert_passed_with_discard_warning(result, caplog, reason): assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] - assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + discard_warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + and "'masker'" in record.getMessage() + and "discarded" in record.getMessage() + ] + assert len(discard_warnings) == 1 + assert reason in discard_warnings[0] + assert "text rewrites included" in discard_warnings[0] assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) @@ -1159,7 +1171,7 @@ async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(translation, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] assert translation.seen_guardrail_names == ["masker"] @@ -1196,7 +1208,7 @@ async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(m with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_TextTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] @@ -1258,7 +1270,9 @@ async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_WritingTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning( + result, caplog, "the guardrail returned 0 tool calls for a stream that carried 1" + ) assert chunks == [_chunk()] @@ -1270,7 +1284,7 @@ async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_ with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_TextTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] @@ -1326,7 +1340,7 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_RefusingTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the translation refused it") assert chunks == [_chunk()] @@ -1561,7 +1575,7 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook returned 2 texts for a stream that carried 1") assert chunks == [_chunk()] @@ -1576,7 +1590,7 @@ async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(m with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") assert chunks == [_chunk()] @@ -1588,7 +1602,9 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls( with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning( + result, caplog, "the legacy hook returned 0 tool calls for a stream that carried 1" + ) assert chunks == [_chunk()] @@ -1620,7 +1636,7 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() ) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") assert chunks == [_tool_only_chunk()] @@ -1658,7 +1674,7 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook's response could not be rescanned") assert chunks == [_chunk()] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 65160a97c54da63f24bf674f4f03551b22ac97c3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:55:03 -0700 Subject: [PATCH 028/349] fix(guardrails): keep tool calls carried by a later choice of a packed multi-choice chunk The rebuild's tool-call selection and its text-only fast path only looked at choice 0 of each chunk, so a chunk that packs several choices (Gemini with candidateCount above 1) lost a tool call carried by a later candidate, and a chunk whose later choice had no tool calls at all made the rebuild raise. Both now consider every choice in the chunk. --- .../streaming_chunk_builder_utils.py | 4 +- litellm/main.py | 66 +++++++++++-------- tests/test_litellm/test_main.py | 62 +++++++++++++++++ 3 files changed, 104 insertions(+), 28 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 5ffe36573d5..f5b723755aa 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -427,7 +427,7 @@ class ChunkProcessor: if not delta: continue choice_index = choice.get("index", 0) - for tool_call in delta.get("tool_calls", ()): + for tool_call in delta.get("tool_calls") or (): if not tool_call: continue if isinstance(tool_call, dict): @@ -478,7 +478,7 @@ class ChunkProcessor: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) - tool_calls = delta.get("tool_calls", []) + tool_calls = delta.get("tool_calls") or () choice_index = choice.get("index", 0) for tool_call in tool_calls: diff --git a/litellm/main.py b/litellm/main.py index 17edafcdfca..a4a648acd4e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8749,6 +8749,39 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o setattr(usage, "cost", computed_cost) +_NON_TEXT_DELTA_FIELDS: Final = ( + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "annotations", + "audio", + "images", + "provider_specific_fields", +) + + +def _stream_choice_delta(choice: object) -> Mapping[str, object]: + delta: Final = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) + if isinstance(delta, Mapping): + return delta + if isinstance(delta, BaseModel): + return delta.model_dump() + return {} + + +def _delta_carries_more_than_text(delta: Mapping[str, object]) -> bool: + return any(delta.get(field) is not None for field in _NON_TEXT_DELTA_FIELDS) + + +def _simple_text_part(choices: Sequence[object]) -> str | None: + deltas: Final = tuple(_stream_choice_delta(choice) for choice in choices) + if any(_delta_carries_more_than_text(delta) for delta in deltas): + return None + content: Final = deltas[0].get("content") + return content if isinstance(content, str) else "" + + def stream_chunk_builder( chunks: list, messages: Sequence | None = None, @@ -8793,31 +8826,11 @@ def stream_chunk_builder( if not chunk.get("choices"): continue - choice = chunk["choices"][0] - delta_obj = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) - if isinstance(delta_obj, dict): - delta = delta_obj - elif hasattr(delta_obj, "model_dump"): - delta = cast(dict[str, Any], delta_obj.model_dump()) - else: - delta = {} - - if ( - delta.get("tool_calls") is not None - or delta.get("function_call") is not None - or delta.get("reasoning_content") is not None - or delta.get("thinking_blocks") is not None - or delta.get("annotations") is not None - or delta.get("audio") is not None - or delta.get("images") is not None - or delta.get("provider_specific_fields") is not None - ): + if (part := _simple_text_part(chunk["choices"])) is None: is_simple_text_stream = False break - - content = delta.get("content") - if isinstance(content, str) and content: - simple_content_parts.append(content) + if part: + simple_content_parts.append(part) if is_simple_text_stream: if simple_content_parts: @@ -8854,9 +8867,10 @@ def stream_chunk_builder( tool_call_chunks: Final = [ chunk for chunk in chunks - if chunk.get("choices") - and "tool_calls" in chunk["choices"][0]["delta"] - and chunk["choices"][0]["delta"]["tool_calls"] is not None + if any( + "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None + for choice in chunk.get("choices") or () + ) ] if len(tool_call_chunks) > 0: diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4f7a51eb531..42599f45ade 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1659,6 +1659,68 @@ async def test_async_mock_delay(): assert delay >= 0.01 +def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): + from litellm import stream_chunk_builder + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-multi-choice", + created=1751934860, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=choices, + ) + + chunks = [ + chunk( + [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), + StreamingChoices( + index=1, + delta=Delta( + role="assistant", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + index=0, + type="function", + function=Function(name="lookup_fruit", arguments='{"fruit":'), + ) + ], + ), + ), + ] + ), + chunk( + [ + StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), + StreamingChoices( + index=1, + delta=Delta( + tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] + ), + finish_reason="tool_calls", + ), + ] + ), + ] + + response = stream_chunk_builder(chunks=chunks) + + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None + assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ + ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') + ] + + def test_stream_chunk_builder_thinking_blocks(): from litellm import stream_chunk_builder from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices From 785c6cffc4826ef44c73a797981e28a294a45668 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Thu, 13 Aug 2026 23:34:40 -0400 Subject: [PATCH 029/349] fix(cost): carry image and video input tokens through the Responses usage bridge Realtime cost is computed from *_tokens_details after the usage round-trips through the Responses shape, and the input half of that shape carried audio only, so image and video prompt tokens stopped being billable as themselves. Vertex splits prompt tokens by modality, so a session sending camera frames arrives with image_tokens set. Those were folded into text_tokens and lost their attribution. The amount happens not to move today, because the calculator falls back to input_cost_per_token when no per-modality rate is set, but the tokens have to survive before any such rate can ever apply. InputTokensDetails now declares image_tokens and video_tokens instead of leaning on pydantic extras, the repeated per-field copying is a loop over the modality names so adding a modality no longer adds a branch, and the read-back in ResponseAPILoggingUtils picks up video_tokens, which PromptTokensDetailsWrapper already declared. The output half of the original change is dropped: 449c091391 landed the same OutputTokensDetails.audio_tokens fix upstream, with its own coverage in test_gemini_realtime_transformation.py, and it always sets output_tokens_details rather than only when non-empty. That structure is kept as upstream wrote it. --- .../transformation.py | 2 ++ litellm/responses/utils.py | 1 + litellm/types/llms/openai.py | 2 ++ .../test_litellm_completion_responses.py | 33 +++++++++++++++++++ 4 files changed, 38 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 01fb6cb483d..ffd7ce491b1 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2851,6 +2851,8 @@ class LiteLLMCompletionResponsesConfig: cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, text_tokens=prompt_details.text_tokens, audio_tokens=prompt_details.audio_tokens, + image_tokens=prompt_details.image_tokens, + video_tokens=prompt_details.video_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 41a3ded7022..f50c17aff85 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1182,6 +1182,7 @@ class ResponseAPILoggingUtils: cached_tokens_details=getattr( response_api_usage.input_tokens_details, "cached_tokens_details", None ), + video_tokens=getattr(response_api_usage.input_tokens_details, "video_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..98548705979 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1291,7 +1291,9 @@ class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 cached_tokens_details: CachedTokensDetails | None = None + image_tokens: int | None = None text_tokens: int | None = None + video_tokens: int | None = None model_config = {"extra": "allow"} diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..2f9f7adfcf1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2885,6 +2885,39 @@ class TestUsageTransformation: assert getattr(response_usage.input_tokens_details, "cache_write_tokens", None) == 800 assert response_usage.input_tokens_details.model_dump()["cache_write_tokens"] == 800 + def test_transform_usage_preserves_input_modality_tokens(self): + """Regression: the bridge dropped image and video input tokens. + + Vertex reports prompt tokens split by modality, so a Live session that sends + camera frames arrives with image_tokens set. InputTokensDetails declared only + audio/cached/text, so those tokens were folded into text and lost their + attribution, and any per-modality rate could never apply to them. + """ + usage = Usage( + prompt_tokens=300, + completion_tokens=10, + total_tokens=310, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, audio_tokens=80, image_tokens=150, video_tokens=50, cached_tokens=0 + ), + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), + ) + + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=usage + ) + details = response_usage.input_tokens_details + assert details is not None + assert getattr(details, "image_tokens", None) == 150 + assert getattr(details, "video_tokens", None) == 50 + assert getattr(details, "audio_tokens", None) == 80 + + from litellm.responses.utils import ResponseAPILoggingUtils + + back = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_usage.model_dump()) + assert back.prompt_tokens_details.image_tokens == 150 + assert back.prompt_tokens_details.video_tokens == 50 + def test_transform_usage_with_reasoning_tokens_gemini(self): """Test that reasoning_tokens from Gemini are properly transformed to output_tokens_details""" # Setup: Simulate Gemini usage with thoughtsTokenCount From 76488beaf8d1a44e0f07b6a4a66c06b9b4390222 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:55:25 -0700 Subject: [PATCH 030/349] fix(utils): reject an untranslatable tool_choice with a 400 instead of a 500 --- litellm/main.py | 2 +- litellm/utils.py | 16 +++- tests/litellm_utils_tests/test_utils.py | 6 +- .../test_validate_tool_choice.py | 74 ++++++++++--------- .../test_litellm_completion_responses.py | 15 ++++ tests/test_litellm/test_main.py | 14 ++++ 6 files changed, 85 insertions(+), 42 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index f6f4ec1bf63..9eea6779abf 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5102,7 +5102,7 @@ def completion( messages = validate_and_fix_openai_messages(messages=messages) tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice - tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice) + tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice, model=model) # validate optional params stop = validate_openai_optional_params(stop=stop) thinking = validate_and_fix_thinking_param(thinking=thinking) diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..17088f0475b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8038,6 +8038,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, + model: str, ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. @@ -8053,12 +8054,19 @@ def validate_chat_completion_tool_choice( # Standard OpenAI format: {"type": "function", "function": {...}} if tool_choice.get("type") is None or tool_choice.get("function") is None: - raise Exception( - f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec" + raise BadRequestError( + message=f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec", + model=model, + llm_provider="", ) return tool_choice - raise Exception( - f"Invalid tool choice, tool_choice={tool_choice}. Got={type(tool_choice)}. Expecting str, or dict. Please ensure tool_choice follows the OpenAI tool_choice spec" + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. Got={type(tool_choice)}. Expecting str, or dict. " + "Please ensure tool_choice follows the OpenAI tool_choice spec" + ), + model=model, + llm_provider="", ) diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 0ccfae55290..11d089719c6 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1334,10 +1334,10 @@ def test_validate_chat_completion_tool_choice(tool_choice, expected_bool): from litellm.utils import validate_chat_completion_tool_choice if expected_bool: - validate_chat_completion_tool_choice(tool_choice=tool_choice) + validate_chat_completion_tool_choice(tool_choice=tool_choice, model="gpt-5.6-sol") else: - with pytest.raises(Exception, match="Invalid tool choice"): - validate_chat_completion_tool_choice(tool_choice=tool_choice) + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"): + validate_chat_completion_tool_choice(tool_choice=tool_choice, model="gpt-5.6-sol") def test_models_by_provider(): diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index b8246fe0deb..b4272af7b90 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -1,60 +1,66 @@ +import re +from typing import Final + import pytest - +import litellm from litellm.utils import validate_chat_completion_tool_choice +MODEL: Final = "anthropic/claude-haiku-4-5" + def test_validate_tool_choice_none(): """Test that None is returned as-is.""" - result = validate_chat_completion_tool_choice(None) + result = validate_chat_completion_tool_choice(None, model=MODEL) assert result is None def test_validate_tool_choice_string(): """Test that string values are returned as-is.""" - assert validate_chat_completion_tool_choice("auto") == "auto" - assert validate_chat_completion_tool_choice("none") == "none" - assert validate_chat_completion_tool_choice("required") == "required" + assert validate_chat_completion_tool_choice("auto", model=MODEL) == "auto" + assert validate_chat_completion_tool_choice("none", model=MODEL) == "none" + assert validate_chat_completion_tool_choice("required", model=MODEL) == "required" def test_validate_tool_choice_standard_dict(): """Test standard OpenAI format with function.""" tool_choice = {"type": "function", "function": {"name": "my_function"}} - result = validate_chat_completion_tool_choice(tool_choice) + result = validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert result == tool_choice def test_validate_tool_choice_cursor_format(): """Cursor IDE format {"type": "auto"} is unwrapped to the bare string.""" - assert validate_chat_completion_tool_choice({"type": "auto"}) == "auto" - assert validate_chat_completion_tool_choice({"type": "none"}) == "none" - assert validate_chat_completion_tool_choice({"type": "required"}) == "required" + assert validate_chat_completion_tool_choice({"type": "auto"}, model=MODEL) == "auto" + assert validate_chat_completion_tool_choice({"type": "none"}, model=MODEL) == "none" + assert validate_chat_completion_tool_choice({"type": "required"}, model=MODEL) == "required" -def test_validate_tool_choice_invalid_dict(): - """Test that invalid dict formats raise exceptions.""" - # Missing both type and function - with pytest.raises(Exception, match='Invalid tool choice, tool_choice=\\{\\}\\. Please ensure') as exc_info: - validate_chat_completion_tool_choice({}) - assert "Invalid tool choice" in str(exc_info.value) - - # Invalid type value - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'invalid'\\}\\.") as exc_info: - validate_chat_completion_tool_choice({"type": "invalid"}) - assert "Invalid tool choice" in str(exc_info.value) - - # Has type but missing function when type is "function" - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'function'\\}\\.") as exc_info: - validate_chat_completion_tool_choice({"type": "function"}) - assert "Invalid tool choice" in str(exc_info.value) +@pytest.mark.parametrize( + "tool_choice", + [ + {}, + {"type": "invalid"}, + {"type": "function"}, + {"name": "lookup_fruit"}, + {"type": "file_search"}, + ], +) +def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): + """A dict shape chat completions cannot carry is the caller's mistake: a 400 that names the field, never a 500.""" + with pytest.raises( + litellm.BadRequestError, match=f"Invalid tool choice, tool_choice={re.escape(str(tool_choice))}\\. Please ensure" + ) as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == MODEL -def test_validate_tool_choice_invalid_type(): - """Test that invalid types raise exceptions.""" - with pytest.raises(Exception, match="\\. Expecting str, or dict\\. Please ensure") as exc_info: - validate_chat_completion_tool_choice(123) - assert "Got=" in str(exc_info.value) - - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\[\\]\\. Got=\\.") as exc_info: - validate_chat_completion_tool_choice([]) - assert "Got=" in str(exc_info.value) +@pytest.mark.parametrize("tool_choice", [123, []]) +def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): + """A non-str, non-dict tool_choice is rejected as a 400 that names the type it got.""" + with pytest.raises( + litellm.BadRequestError, match=f"Got={re.escape(str(type(tool_choice)))}\\. Expecting str, or dict\\." + ) as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..3950fd549ef 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4906,3 +4906,18 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +@pytest.mark.parametrize("stream", [True, False]) +async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool): + with pytest.raises(litellm.BadRequestError) as exc_info: + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Which fruit is red?", + tools=[{"type": "function", "name": "lookup_fruit", "parameters": {"type": "object"}}], + tool_choice={"type": "file_search"}, + stream=stream, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3dccb2b35bf..81ad161772e 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3850,3 +3850,17 @@ def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_ assert "extra_headers" not in body assert body["model"] == "gpt-5.4" assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS + + +@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) +def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Which fruit is red?"}], + tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert f"tool_choice={tool_choice}" in str(exc_info.value) From 2bbf34c6520ef3266a62121510470868f73e499e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:00:33 -0700 Subject: [PATCH 031/349] fix(utils): keep the tool_choice validator's model argument optional --- litellm/utils.py | 2 +- tests/litellm_utils_tests/test_validate_tool_choice.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 17088f0475b..c2dbbda2b68 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8038,7 +8038,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, - model: str, + model: str = "", ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index b4272af7b90..a9dacf9fa15 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -64,3 +64,11 @@ def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): ) as exc_info: validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert exc_info.value.status_code == 400 + + +def test_validate_tool_choice_without_model_is_still_a_400(): + """Callers that predate the model argument keep getting a 400, with an empty model on the error.""" + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc_info: + validate_chat_completion_tool_choice({"type": "bogus"}) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == "" From 61ac4f57394e022382459235baa598ac89ce00d1 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:03:56 +0000 Subject: [PATCH 032/349] test(e2e): add scripted-provider cost calculation suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 3 +- tests/e2e/conftest.py | 9 +- tests/e2e/cost_calculation/conftest.py | 139 ++++ tests/e2e/cost_calculation/cost_matrix.py | 458 +++++++++++++ tests/e2e/cost_calculation/scripted_client.py | 70 ++ .../e2e/cost_calculation/scripted_provider.py | 631 ++++++++++++++++++ .../test_token_pricing_e2e.py | 115 ++++ .../cost_calculation/test_wire_formats_e2e.py | 186 ++++++ tests/e2e/cost_map.json | 352 ++++++++++ .../coverage_registry/quota_management.yaml | 2 + tests/e2e/e2e_config.py | 16 + tests/e2e/pytest.ini | 1 + 12 files changed, 1980 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/cost_calculation/conftest.py create mode 100644 tests/e2e/cost_calculation/cost_matrix.py create mode 100644 tests/e2e/cost_calculation/scripted_client.py create mode 100644 tests/e2e/cost_calculation/scripted_provider.py create mode 100644 tests/e2e/cost_calculation/test_token_pricing_e2e.py create mode 100644 tests/e2e/cost_calculation/test_wire_formats_e2e.py create mode 100644 tests/e2e/cost_map.json diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0541ce25d4b..b6c3840f626 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,6 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` @@ -221,7 +222,7 @@ other... ``` ## Hard Rules -- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description +- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests; the one carve-out is a scripted upstream served through a real HTTP sidecar (the cost_calculation suite's scripted provider), allowed because provider-response-shape coverage needs a controlled usage payload and every hop from the proxy's upstream call to the spend row still executes for real. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description - use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index b1a75d5f862..7ab41b8ff68 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -22,9 +22,9 @@ from typing import Final import pytest import requests - from e2e_config import ( CONTROL_PLANE_BASE_URL, + COST_MAP_OPT_IN_ENV, FIXTURE_DIR, FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, @@ -53,6 +53,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "cost_map_stack": COST_MAP_OPT_IN_ENV, } ) @@ -120,6 +121,12 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set", ) + config.addinivalue_line( + "markers", + "cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json " + "(LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless " + "E2E_COST_MAP_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py new file mode 100644 index 00000000000..1bba3d50e1d --- /dev/null +++ b/tests/e2e/cost_calculation/conftest.py @@ -0,0 +1,139 @@ +"""Cost-calculation suite fixtures. + +Runs against a dedicated proxy whose whole model cost map is the test-owned +``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL), so every deployment +bills at rates the test asserts literal arithmetic on. Provider calls are +answered by the scripted-provider sidecar (``scripted_provider.py``), registered +per scenario over its control API. + +Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). +""" + +from __future__ import annotations + +import importlib.util +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Final, Protocol, cast + +import pytest + +from cost_matrix import Case, FrontierModel +from e2e_config import COST_MAP_PROXY_URL +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody +from proxy_client import ProxyClient, build_proxy_client +from scripted_client import ScenarioHandle, delete_scenario, register_scenario +from scripted_provider import Scenario + + +def _load_cost_rows() -> ModuleType: + """Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree + has no package layout), the same trick the mcp suite uses for + logging/datadog_reader.py.""" + path = ( + Path(__file__).resolve().parent.parent + / "quota_management" + / "spend_tracking" + / "cost_rows.py" + ) + name = "e2e_spend_tracking_cost_rows" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +class SpendCostBreakdown(Protocol): + input_cost: float | None + output_cost: float | None + cache_read_cost: float | None + cache_creation_cost: float | None + reasoning_cost: float | None + tool_usage_cost: float | None + total_cost: float | None + service_tier: str | None + + def model_dump(self) -> dict[str, object]: ... + + +class SpendRowMetadata(Protocol): + cost_breakdown: SpendCostBreakdown | None + + +class SpendCostRow(Protocol): + """The slice of spend_tracking.cost_rows.CostRow this suite reads.""" + + spend: float | None + prompt_tokens: int | None + completion_tokens: int | None + metadata: SpendRowMetadata | None + + @property + def breakdown(self) -> SpendCostBreakdown: ... + + +class CostRowsModule(Protocol): + """cost_rows.py loaded by path has no importable name for basedpyright, so + its surface is declared here and reached through a single cast.""" + + approx_equal: Callable[[float, float], bool] + assert_total_is_sum_of_components: Callable[[SpendCostRow], None] + poll_cost_row_where: Callable[ + [ProxyClient, str, Callable[[SpendCostRow], bool]], SpendCostRow | None + ] + + +cost_rows: Final[CostRowsModule] = cast(CostRowsModule, _load_cost_rows()) + + +@dataclass(frozen=True, slots=True) +class CostCalcClient: + """The suite's client: a ProxyClient pointed at the cost-map proxy pod.""" + + proxy: ProxyClient + + +@pytest.fixture(scope="session") +def client() -> CostCalcClient: + proxy = build_proxy_client( + base_url=COST_MAP_PROXY_URL, + control_plane_base_url=COST_MAP_PROXY_URL, + replica_urls=(COST_MAP_PROXY_URL,), + ) + return CostCalcClient(proxy=proxy) + + +def register_scenario_deployment( + client: CostCalcClient, + resources: ResourceManager, + model: FrontierModel, + case: Case, + marker: str, +) -> tuple[str, ScenarioHandle]: + """Register the case's scenario on the sidecar plus a deployment pointed at + it; both are torn down by ``resources``. Returns the callable model_name.""" + scenario: Scenario = case.scenario( + scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}" + ) + handle = register_scenario(scenario) + resources.defer(lambda: delete_scenario(handle)) + model_name = f"{model.model_name}-{marker}" + model_id = client.proxy.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody( + model=model.litellm_model, + api_key="sk-scripted-provider", + api_base=handle.api_base(), + ), + model_info=ModelInfoBody(), + ) + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model_name, handle diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py new file mode 100644 index 00000000000..bc466d7d823 --- /dev/null +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -0,0 +1,458 @@ +"""The cost-calculation matrix: frontier model set, the pricing-component cases +each model runs, and the expected-cost arithmetic. + +Rates come from ``tests/e2e/cost_map.json``, which the proxy under test loads as +its ENTIRE model cost map (LITELLM_MODEL_COST_MAP_URL), so an entry's rates are +exactly what the proxy bills and nothing in the suite depends on the bundled +map. Each model's rates are a distinct multiple of a shared base set, so a +component billed at the wrong model's rate (or the wrong case's rate) can never +coincidentally match. + +Case applicability is pricing-field-gated AND wire-gated: a case runs for a +model only when the entry carries the rate the case exercises and the wire can +report the token kind that rate prices. When the wire cannot report a kind +(e.g. Anthropic has no reasoning-token field, Responses reports no cache +creation), the case is absent from the matrix rather than silently zero. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from scripted_provider import Scenario, ScriptedOutput, ScriptedUsage, Wire + +COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" + + +class SearchContextCostPerQuery(BaseModel): + model_config = ConfigDict(frozen=True) + + search_context_size_low: float | None = None + search_context_size_medium: float | None = None + search_context_size_high: float | None = None + + +class CostMapEntry(BaseModel): + """The pricing fields of a cost-map entry the matrix reads. Shaped like a + ``model_prices_and_context_window.json`` entry; unmodelled keys are ignored.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_provider: str + mode: str + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + cache_creation_input_token_cost_above_1hr: float | None = None + output_cost_per_reasoning_token: float | None = None + input_cost_per_audio_token: float | None = None + output_cost_per_audio_token: float | None = None + input_cost_per_token_above_200k_tokens: float | None = None + output_cost_per_token_above_200k_tokens: float | None = None + input_cost_per_token_flex: float | None = None + output_cost_per_token_flex: float | None = None + input_cost_per_token_priority: float | None = None + output_cost_per_token_priority: float | None = None + search_context_cost_per_query: SearchContextCostPerQuery | None = None + web_search_billing_unit: str | None = None + + +_COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) +_COST_MAP: Final[dict[str, CostMapEntry]] = _COST_MAP_ADAPTER.validate_python( + json.loads(COST_MAP_PATH.read_text()) +) + +TIER_THRESHOLD_TOKENS: Final = 200_000 + + +@dataclass(frozen=True, slots=True) +class FrontierModel: + """One deployment under test: the model_name the suite registers, the + provider-prefixed litellm model string, the wire the scripted upstream + speaks, its cost-map key, and the sibling map model the response_model + override case reports.""" + + model_name: str + litellm_model: str + wire: Wire + map_key: str + override_model: str + + @property + def rates(self) -> CostMapEntry: + return _COST_MAP[self.map_key] + + @property + def override_rates(self) -> CostMapEntry: + return _COST_MAP[self.override_map_key] + + @property + def override_map_key(self) -> str: + return _OVERRIDE_MAP_KEYS[self.override_model] + + @property + def provider(self) -> str: + return self.rates.litellm_provider + + @property + def api_key(self) -> str: + # The scripted upstream ignores auth; a fixed bogus key proves the suite + # spends zero real provider calls. + return "sk-scripted-provider" + + +# Response-model override targets: emit a sibling's bare provider-facing name so +# the biller's provider-prefixed lookup lands on that sibling's map key. +_OVERRIDE_MODELS: Final[dict[str, str]] = { + "gpt-5.6": "gpt-5.4-mini", + "gpt-5.5-pro": "gpt-5.3-codex", + "gpt-5.3-codex": "gpt-5.5-pro", + "gpt-5.4-mini": "gpt-5.6", + "claude-opus-5": "claude-sonnet-5", + "claude-sonnet-5": "claude-opus-5", + "claude-haiku-4-5": "claude-sonnet-5", + "gemini/gemini-3.8-flash": "gemini-3.1-pro-preview", + "gemini/gemini-3.1-pro-preview": "gemini-3.8-flash", + "together_ai/moonshotai/Kimi-K3": "zai-org/GLM-5.3", + "together_ai/zai-org/GLM-5.3": "moonshotai/Kimi-K3", + "fireworks_ai/kimi-k3": "qwen3p8-max", + "fireworks_ai/qwen3p8-max": "kimi-k3", + "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", +} + +_OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { + "gpt-5.4-mini": "gpt-5.4-mini", + "gpt-5.6": "gpt-5.6", + "gpt-5.3-codex": "gpt-5.3-codex", + "gpt-5.5-pro": "gpt-5.5-pro", + "claude-sonnet-5": "claude-sonnet-5", + "claude-opus-5": "claude-opus-5", + "gemini-3.1-pro-preview": "gemini/gemini-3.1-pro-preview", + "gemini-3.8-flash": "gemini/gemini-3.8-flash", + "zai-org/GLM-5.3": "together_ai/zai-org/GLM-5.3", + "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", + "qwen3p8-max": "fireworks_ai/qwen3p8-max", + "kimi-k3": "fireworks_ai/kimi-k3", +} + + +_FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( + ("gpt-5.6", "openai/gpt-5.6", "openai_chat"), + ("gpt-5.5-pro", "openai/gpt-5.5-pro", "openai_responses"), + ("gpt-5.3-codex", "openai/gpt-5.3-codex", "openai_responses"), + ("gpt-5.4-mini", "openai/gpt-5.4-mini", "openai_chat"), + ("claude-opus-5", "anthropic/claude-opus-5", "anthropic_messages"), + ("claude-sonnet-5", "anthropic/claude-sonnet-5", "anthropic_messages"), + ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "anthropic_messages"), + ("gemini/gemini-3.8-flash", "gemini/gemini-3.8-flash", "gemini_generate"), + ("gemini/gemini-3.1-pro-preview", "gemini/gemini-3.1-pro-preview", "gemini_generate"), + ("together_ai/moonshotai/Kimi-K3", "together_ai/moonshotai/Kimi-K3", "together_chat"), + ("together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3", "together_chat"), + ("fireworks_ai/kimi-k3", "fireworks_ai/kimi-k3", "fireworks_chat"), + ("fireworks_ai/qwen3p8-max", "fireworks_ai/qwen3p8-max", "fireworks_chat"), + ("fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/deepseek-v4p1-flash", "fireworks_chat"), +) + + +def _frontier() -> tuple[FrontierModel, ...]: + return tuple( + FrontierModel( + model_name=f"cc-{map_key.replace('/', '-').lower()}", + litellm_model=litellm_model, + wire=wire, + map_key=map_key, + override_model=_OVERRIDE_MODELS[map_key], + ) + for map_key, litellm_model, wire in _FRONTIER_SPECS + ) + + +FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() + +# Token kinds each wire can report, gating which pricing cases apply. +_WIRE_CAPS: Final[dict[str, frozenset[str]]] = { + "openai_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), + "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), + # Product gap: litellm hard-indexes message_delta["usage"] in + # anthropic/chat/handler.py, so a usage-absent anthropic stream raises + # KeyError; the real wire always carries it, so the case cannot be + # represented. + "anthropic_messages": frozenset({"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model"}), + # Product gap: the gemini transform sets ModelResponse.model from the + # request and drops the provider's modelVersion, so a response-model + # override can never be priced on this wire. + "gemini_generate": frozenset({"cache_read", "reasoning", "audio", "web_search", "absent_usage"}), + "together_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), + "fireworks_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), +} + +CaseName = Literal[ + "basic", + "cache_read", + "cache_write_5m", + "cache_write_1h", + "reasoning", + "audio", + "tiered", + "service_tier_flex", + "service_tier_priority", + "web_search", + "stream", + "stream_no_usage", + "response_model_override", +] + + +@dataclass(frozen=True, slots=True) +class Case: + name: CaseName + usage: ScriptedUsage + stream: bool = False + stream_usage: Literal["final_chunk", "absent"] = "final_chunk" + service_tier: Literal["flex", "priority"] | None = None + # For web_search the wire's reported call count is not always what gets + # billed: chat-completions surfaces only expose url_citation annotations, so + # the biller floors to one call; responses/messages/gemini report a real + # count. + billed_web_search_calls: int = 0 + response_model_override: bool = False + exact_spend: bool = True + # stream_usage=absent on a wire with no proxy-side token recount means the + # bill is exactly zero; asserted as such rather than skipped. + expect_zero_bill: bool = False + + def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: + return Scenario( + scenario_id=scenario_id, + wire=model.wire, + usage=self.usage, + output=ScriptedOutput( + text=text, + response_model=model.override_model if self.response_model_override else None, + ), + stream_usage=self.stream_usage, + service_tier=self.service_tier, + ) + + +_BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) + + +def _web_search_case(model: FrontierModel) -> Case: + counts_exactly = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + return Case( + name="web_search", + usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), + billed_web_search_calls=3 if counts_exactly else 1, + ) + + +def cases_for(model: FrontierModel) -> tuple[Case, ...]: + rates = model.rates + caps = _WIRE_CAPS[model.wire] + cases: list[Case] = [Case(name="basic", usage=_BASIC_USAGE)] + if rates.cache_read_input_token_cost is not None and "cache_read" in caps: + cases.append( + Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) + ) + if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps: + cases.append( + Case( + name="cache_write_5m", + usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), + ) + ) + if ( + rates.cache_creation_input_token_cost_above_1hr is not None + and rates.cache_creation_input_token_cost is not None + and "cache_write_1h" in caps + ): + cases.append( + Case( + name="cache_write_1h", + usage=ScriptedUsage( + fresh_input_tokens=90, + cache_write_5m_tokens=20, + cache_write_1h_tokens=40, + output_tokens=30, + ), + ) + ) + if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps: + cases.append( + Case( + name="reasoning", + usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), + ) + ) + if ( + rates.input_cost_per_audio_token is not None + and rates.output_cost_per_audio_token is not None + and "audio" in caps + ): + cases.append( + Case( + name="audio", + usage=ScriptedUsage( + fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 + ), + ) + ) + if ( + rates.input_cost_per_token_above_200k_tokens is not None + and rates.output_cost_per_token_above_200k_tokens is not None + ): + cases.append( + Case( + name="tiered", + usage=ScriptedUsage( + fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 + ), + ) + ) + if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None: + cases.append( + Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") + ) + if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None: + cases.append( + Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") + ) + if rates.search_context_cost_per_query is not None and "web_search" in caps: + cases.append(_web_search_case(model)) + cases.append(Case(name="stream", usage=_BASIC_USAGE, stream=True)) + if "absent_usage" in caps: + cases.append( + Case( + name="stream_no_usage", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + exact_spend=False, + # The responses surface bills only provider-reported usage; + # with no usage in the stream the spend row is zero. Other + # wires recount tokens proxy-side and bill a nonzero amount. + expect_zero_bill=model.wire == "openai_responses", + ) + ) + if "response_model" in caps: + cases.append(Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True)) + return tuple(cases) + + +@dataclass(frozen=True, slots=True) +class ExpectedCost: + """The expected bill split the way the spend row's cost_breakdown reports + it: the gross input component (cache reads/writes folded in), the output + component, and the tool-usage component.""" + + input_cost: float + output_cost: float + tool_cost: float + + @property + def total(self) -> float: + return self.input_cost + self.output_cost + self.tool_cost + + +def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: + """Literal arithmetic on the test-map rates over the scripted token counts. + + Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; + output = text*out + reasoning*reasoning + audio_out*audio_out; plus the + billed web-search calls at the medium search-context rate. Above-threshold + swaps every input/output rate to its ``_above_200k_tokens`` variant when + total prompt tokens exceed the threshold; a service tier swaps input/output + to the tier's variants, falling back to the base rate when a variant is + unset -- mirroring _get_token_base_cost in litellm's cost calculator. + """ + rates = model.override_rates if case.response_model_override else model.rates + u = case.usage + prompt_tokens = ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + u.audio_input_tokens + ) + tiered = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate = rates.input_cost_per_token or 0.0 + out_rate = rates.output_cost_per_token or 0.0 + if case.service_tier == "flex": + in_rate = rates.input_cost_per_token_flex or in_rate + out_rate = rates.output_cost_per_token_flex or out_rate + if case.service_tier == "priority": + in_rate = rates.input_cost_per_token_priority or in_rate + out_rate = rates.output_cost_per_token_priority or out_rate + if tiered: + in_rate = rates.input_cost_per_token_above_200k_tokens or in_rate + out_rate = rates.output_cost_per_token_above_200k_tokens or out_rate + input_cost = ( + u.fresh_input_tokens * in_rate + + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + ) + output_cost = ( + u.output_tokens * out_rate + + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + ) + search = rates.search_context_cost_per_query + tool_cost = case.billed_web_search_calls * ( + search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + ) + return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) + + +def expected_cost(model: FrontierModel, case: Case) -> float: + return expected_breakdown(model, case).total + + +def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: + """(prompt_tokens, completion_tokens) the spend row should carry, per the + wire's normalization: Anthropic folds cache read/write into prompt_tokens, + everyone else reports the totals the wire emitted.""" + u = case.usage + if model.wire == "anthropic_messages": + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, + u.output_tokens, + ) + if model.wire == "gemini_generate": + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + if model.wire == "openai_responses": + return ( + u.fresh_input_tokens + u.cache_read_tokens, + u.output_tokens + u.reasoning_tokens, + ) + return ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) diff --git a/tests/e2e/cost_calculation/scripted_client.py b/tests/e2e/cost_calculation/scripted_client.py new file mode 100644 index 00000000000..dceec02630a --- /dev/null +++ b/tests/e2e/cost_calculation/scripted_client.py @@ -0,0 +1,70 @@ +"""Client side of the scripted-provider sidecar: register scenarios over its +control API through the shared transport helpers and get back a handle whose +``api_base`` is what a /model/new deployment should register for the proxy to +reach the scripted wire.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BASE +from e2e_http import URL, NoBody, unwrap, post +from e2e_http import delete as http_delete +from scripted_provider import ( + Scenario, + ScenarioDeleted, + ScenarioRegistered, + Wire, +) + + +@dataclass(frozen=True, slots=True) +class ScenarioHandle: + scenario_id: str + wire: Wire + proxy_base: str + + def api_base(self) -> str: + return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}" + + def _mount(self) -> str: + return { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", + }[self.wire] + + +def register_scenario(scenario: Scenario) -> ScenarioHandle: + """POST the scenario to the sidecar's control API and return its handle.""" + result = unwrap( + post( + URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"), + headers=NoBody(), + json=scenario, + response_type=ScenarioRegistered, + ) + ) + return ScenarioHandle( + scenario_id=result.scenario_id, + wire=scenario.wire, + proxy_base=SCRIPTED_PROVIDER_PROXY_BASE, + ) + + +def delete_scenario(handle: ScenarioHandle) -> None: + unwrap( + http_delete( + URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios/{handle.scenario_id}"), + headers=NoBody(), + json=NoBody(), + response_type=ScenarioDeleted, + ) + ) + + +CONTROL_URL: Final = SCRIPTED_PROVIDER_CONTROL_URL diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py new file mode 100644 index 00000000000..93a6f49ec25 --- /dev/null +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -0,0 +1,631 @@ +"""Scripted provider sidecar for the cost-calculation e2e suite. + +A standalone process (``python -m cost_calculation.scripted_provider``) that +pretends to be an LLM provider for the proxy under test. The suite registers a +Scenario over a small control API; the provider wire routes then answer the +proxy's upstream calls with the scripted usage figures, in the exact wire shape +the real provider would emit (OpenAI chat completions, OpenAI Responses, +Anthropic Messages, Gemini generateContent, or the OpenAI-compatible Together / +Fireworks surfaces). Because the usage is scripted, expected spend is literal +arithmetic on the test cost map's rates, with no dependency on what a real +provider would report. + +Layout on one port: + +- ``GET /health`` liveness +- ``POST /_scenarios`` register a Scenario JSON, returns its id +- ``DELETE /_scenarios/`` remove it +- ``POST ///`` provider wire; mount is one of + ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks`` and the + remainder is whatever path the provider client appends (``chat/completions``, + ``responses``, ``v1/messages``, ``models/:generateContent`` ...) + +A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini +verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the +final stream chunk carries usage or the provider reports none. +""" + +from __future__ import annotations + +import json +import sys +import threading +import time +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, Literal +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +Wire = Literal[ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate", + "together_chat", + "fireworks_chat", +] + +_WIRE_MOUNTS: Final[dict[str, str]] = { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", +} + +StreamUsage = Literal["final_chunk", "absent"] +ServiceTier = Literal["flex", "priority"] + + +class ScriptedUsage(BaseModel): + """Physical token counts the scripted response reports. ``fresh_input_tokens`` + is the uncached, never-written, non-audio input count; ``output_tokens`` is + the non-reasoning, non-audio output count. Renderers add the cached, written, + audio, and reasoning counts into the wire's total fields the way the real + provider does (inside prompt_tokens for OpenAI/Gemini, as uncached-only + input_tokens for Anthropic).""" + + model_config = ConfigDict(frozen=True) + + fresh_input_tokens: int = 0 + output_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_5m_tokens: int = 0 + cache_write_1h_tokens: int = 0 + reasoning_tokens: int = 0 + audio_input_tokens: int = 0 + audio_output_tokens: int = 0 + web_search_calls: int = 0 + + +class ScriptedOutput(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str + finish_reason: str = "stop" + # When set, emitted verbatim as the response's model field, letting a test + # prove the biller prices the provider-reported model. + response_model: str | None = None + # OpenAI-compatible providers can report a provider-computed cost; emitted as + # the top-level "cost" field on the together/fireworks wire. + provider_cost: float | None = None + + +class Scenario(BaseModel): + model_config = ConfigDict(frozen=True) + + scenario_id: str + wire: Wire + usage: ScriptedUsage + output: ScriptedOutput + stream_usage: StreamUsage = "final_chunk" + service_tier: ServiceTier | None = None + + @property + def mount(self) -> str: + return _WIRE_MOUNTS[self.wire] + + +class ScenarioRegistered(BaseModel): + scenario_id: str + + +class ScenarioDeleted(BaseModel): + deleted: bool + + +class HealthStatus(BaseModel): + status: str + + +@dataclass(frozen=True, slots=True) +class RenderedResponse: + status_code: int + content_type: str + body: bytes + + +def _json_bytes(payload: dict[str, object]) -> bytes: + return json.dumps(payload).encode("utf-8") + + +def _sse(events: tuple[tuple[str | None, dict[str, object] | str], ...]) -> bytes: + frames: list[str] = [] + for event_name, data in events: + head = f"event: {event_name}\n" if event_name is not None else "" + payload = data if isinstance(data, str) else json.dumps(data) + frames.append(f"{head}data: {payload}\n\n") + return "".join(frames).encode("utf-8") + + +# ---------- per-wire usage shapes ---------- + + +def _openai_usage(u: ScriptedUsage) -> dict[str, object]: + prompt_tokens = ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens + ) + completion_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + prompt_details: dict[str, object] = {} + if u.cache_read_tokens: + prompt_details["cached_tokens"] = u.cache_read_tokens + if u.cache_write_5m_tokens or u.cache_write_1h_tokens: + prompt_details["cache_write_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens + prompt_details["cache_creation_token_details"] = { + "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, + "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, + } + if u.audio_input_tokens: + prompt_details["audio_tokens"] = u.audio_input_tokens + completion_details: dict[str, object] = {} + if u.reasoning_tokens: + completion_details["reasoning_tokens"] = u.reasoning_tokens + if u.audio_output_tokens: + completion_details["audio_tokens"] = u.audio_output_tokens + usage: dict[str, object] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + if prompt_details: + usage["prompt_tokens_details"] = prompt_details + if completion_details: + usage["completion_tokens_details"] = completion_details + return usage + + +def _anthropic_usage(u: ScriptedUsage) -> dict[str, object]: + # Anthropic reports uncached-only input_tokens; cache reads and writes ride + # top-level fields, with the 5m/1h write split under cache_creation. + usage: dict[str, object] = { + "input_tokens": u.fresh_input_tokens, + "output_tokens": u.output_tokens, + } + if u.cache_read_tokens: + usage["cache_read_input_tokens"] = u.cache_read_tokens + if u.cache_write_5m_tokens or u.cache_write_1h_tokens: + usage["cache_creation_input_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens + usage["cache_creation"] = { + "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, + "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, + } + if u.web_search_calls: + usage["server_tool_use"] = {"web_search_requests": u.web_search_calls} + return usage + + +def _gemini_usage(u: ScriptedUsage) -> dict[str, object]: + # promptTokenCount carries the cached count inside it; TEXT modality is the + # cached-inclusive text count so litellm's implicit-caching subtraction lands + # on the fresh figure. candidatesTokenCount includes reasoning + audio. + prompt_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + candidates = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + usage: dict[str, object] = { + "promptTokenCount": prompt_tokens, + "candidatesTokenCount": candidates, + "totalTokenCount": prompt_tokens + candidates, + } + if u.cache_read_tokens: + usage["cachedContentTokenCount"] = u.cache_read_tokens + if u.reasoning_tokens: + usage["thoughtsTokenCount"] = u.reasoning_tokens + prompt_details = [{"modality": "TEXT", "tokenCount": u.fresh_input_tokens + u.cache_read_tokens}] + if u.audio_input_tokens: + prompt_details.append({"modality": "AUDIO", "tokenCount": u.audio_input_tokens}) + usage["promptTokensDetails"] = prompt_details + if u.audio_output_tokens: + usage["candidatesTokensDetails"] = [ + {"modality": "TEXT", "tokenCount": u.output_tokens + u.reasoning_tokens}, + {"modality": "AUDIO", "tokenCount": u.audio_output_tokens}, + ] + return usage + + +def _responses_usage(u: ScriptedUsage) -> dict[str, object]: + input_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + output_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + usage: dict[str, object] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + input_details: dict[str, object] = {} + if u.cache_read_tokens: + input_details["cached_tokens"] = u.cache_read_tokens + if input_details: + usage["input_tokens_details"] = input_details + if u.reasoning_tokens: + usage["output_tokens_details"] = {"reasoning_tokens": u.reasoning_tokens} + return usage + + +# ---------- per-wire responses ---------- + + +def _openai_message(scenario: Scenario) -> dict[str, object]: + message: dict[str, object] = {"role": "assistant", "content": scenario.output.text} + if scenario.usage.web_search_calls: + message["annotations"] = [ + { + "type": "url_citation", + "url_citation": { + "url": "https://scripted.example/source", + "title": "scripted source", + "start_index": 0, + "end_index": 1, + }, + } + for _ in range(scenario.usage.web_search_calls) + ] + return message + + +def _openai_chat_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + body: dict[str, object] = { + "id": f"chatcmpl-{scenario.scenario_id}", + "object": "chat.completion", + "created": int(time.time()), + "model": scenario.output.response_model or requested_model, + "choices": [ + { + "index": 0, + "message": _openai_message(scenario), + "finish_reason": scenario.output.finish_reason, + } + ], + "usage": _openai_usage(scenario.usage), + } + if scenario.service_tier is not None: + body["service_tier"] = scenario.service_tier + if scenario.output.provider_cost is not None: + body["cost"] = scenario.output.provider_cost + return body + + +def _openai_chunk(scenario: Scenario, requested_model: str, **kw: object) -> dict[str, object]: + chunk: dict[str, object] = { + "id": f"chatcmpl-{scenario.scenario_id}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": scenario.output.response_model or requested_model, + } + chunk.update(kw) + return chunk + + +def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: + _EMPTY_DELTA: Final[dict[str, object]] = {} + delta: dict[str, object] = {"role": "assistant", "content": scenario.output.text} + if scenario.usage.web_search_calls: + delta["annotations"] = _openai_message(scenario)["annotations"] + events: list[tuple[str | None, dict[str, object] | str]] = [ + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + ), + ), + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ), + ), + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[ + { + "index": 0, + "delta": _EMPTY_DELTA, + "finish_reason": scenario.output.finish_reason, + } + ], + ), + ), + ] + if scenario.stream_usage == "final_chunk": + events.append( + (None, _openai_chunk(scenario, requested_model, choices=(), usage=_openai_usage(scenario.usage))) + ) + events.append((None, "[DONE]")) + return _sse(tuple(events)) + + +def _anthropic_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + return { + "id": f"msg_{scenario.scenario_id}", + "type": "message", + "role": "assistant", + "model": scenario.output.response_model or requested_model, + "content": [{"type": "text", "text": scenario.output.text}], + "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + "usage": _anthropic_usage(scenario.usage), + } + + +def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: + emit_usage = scenario.stream_usage == "final_chunk" + input_usage = {k: v for k, v in _anthropic_usage(scenario.usage).items() if k != "output_tokens"} + message_start: dict[str, object] = { + "type": "message_start", + "message": { + "id": f"msg_{scenario.scenario_id}", + "type": "message", + "role": "assistant", + "model": scenario.output.response_model or requested_model, + "content": [], + "stop_reason": None, + **({"usage": input_usage} if emit_usage else {}), + }, + } + message_delta: dict[str, object] = { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + }, + **({"usage": {"output_tokens": scenario.usage.output_tokens}} if emit_usage else {}), + } + return _sse( + ( + ("message_start", message_start), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": scenario.output.text}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", message_delta), + ("message_stop", {"type": "message_stop"}), + ) + ) + + +def _gemini_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + candidate: dict[str, object] = { + "content": {"parts": [{"text": scenario.output.text}], "role": "model"}, + "finishReason": "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), + "index": 0, + } + if scenario.usage.web_search_calls: + candidate["groundingMetadata"] = { + "webSearchQueries": [f"query {i}" for i in range(scenario.usage.web_search_calls)] + } + return { + "candidates": [candidate], + "usageMetadata": _gemini_usage(scenario.usage), + "modelVersion": scenario.output.response_model or requested_model, + } + + +def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: + first = _gemini_body(scenario, requested_model) + if scenario.stream_usage == "absent": + first = {k: v for k, v in first.items() if k != "usageMetadata"} + events: list[tuple[str | None, dict[str, object] | str]] = [(None, first)] + if scenario.stream_usage == "final_chunk": + events.append( + ( + None, + { + "candidates": [], + "usageMetadata": _gemini_usage(scenario.usage), + "modelVersion": scenario.output.response_model or requested_model, + }, + ) + ) + return _sse(tuple(events)) + + +def _responses_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + output: list[dict[str, object]] = [ + {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} + for i in range(scenario.usage.web_search_calls) + ] + output.append( + { + "type": "message", + "id": f"msg_{scenario.scenario_id}", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": scenario.output.text, + "annotations": [], + } + ], + } + ) + return { + "id": f"resp_{scenario.scenario_id}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": scenario.output.response_model or requested_model, + "output": output, + "usage": _responses_usage(scenario.usage), + } + + +def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: + completed = _responses_body(scenario, requested_model) + if scenario.stream_usage == "absent": + completed = {k: v for k, v in completed.items() if k != "usage"} + created = {**completed, "status": "in_progress", "usage": None} + return _sse( + ( + ("response.created", {"type": "response.created", "response": created}), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": f"msg_{scenario.scenario_id}", + "output_index": scenario.usage.web_search_calls, + "content_index": 0, + "delta": scenario.output.text, + }, + ), + ("response.completed", {"type": "response.completed", "response": completed}), + ) + ) + + +def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> RenderedResponse: + if scenario.wire == "anthropic_messages": + if stream: + return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_anthropic_body(scenario, requested_model))) + if scenario.wire == "gemini_generate": + if stream: + return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model))) + if scenario.wire == "openai_responses": + if stream: + return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model))) + # openai_chat, together_chat, fireworks_chat share the OpenAI chat shape. + if stream: + return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model))) + + +# ---------- registry + request routing ---------- + + +class _ScenarioStore: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self._scenarios: dict[str, Scenario] = {} # mutable-ok: server state, guarded by _lock + + def put(self, scenario: Scenario) -> None: + with self._lock: + self._scenarios[scenario.scenario_id] = scenario + + def drop(self, scenario_id: str) -> bool: + with self._lock: + return self._scenarios.pop(scenario_id, None) is not None + + def get(self, scenario_id: str) -> Scenario | None: + with self._lock: + return self._scenarios.get(scenario_id) + + +_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +def _request_body(body: bytes) -> dict[str, object]: + try: + return _REQUEST_BODY.validate_json(body) + except ValueError: + return {} + + +def _request_wants_stream(path_tail: str, body: bytes) -> bool: + if ":streamGenerateContent" in path_tail: + return True + if not body: + return False + return _request_body(body).get("stream") is True + + +def _request_model(body: bytes) -> str: + model = _request_body(body).get("model") + return model if isinstance(model, str) else "unknown" + + +def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: + path = urlsplit(raw_path).path + segments = [segment for segment in path.split("/") if segment] + if method == "GET" and segments == ["health"]: + return RenderedResponse(200, "application/json", _json_bytes({"status": "ok"})) + if segments and segments[0] == "_scenarios": + if method == "POST" and len(segments) == 1: + try: + scenario = Scenario.model_validate_json(body) + except ValidationError as exc: + return RenderedResponse(400, "application/json", _json_bytes({"error": str(exc)})) + store.put(scenario) + return RenderedResponse(200, "application/json", _json_bytes({"scenario_id": scenario.scenario_id})) + if method == "DELETE" and len(segments) == 2: + deleted = store.drop(segments[1]) + return RenderedResponse( + 200 if deleted else 404, "application/json", _json_bytes({"deleted": deleted}) + ) + return RenderedResponse(404, "application/json", _json_bytes({"error": "unknown control route"})) + if len(segments) < 2 or method != "POST": + return RenderedResponse(404, "application/json", _json_bytes({"error": f"no route for {method} {path}"})) + scenario_id, mount = segments[0], segments[1] + scenario = store.get(scenario_id) + if scenario is None: + return RenderedResponse(404, "application/json", _json_bytes({"error": f"unknown scenario {scenario_id}"})) + if scenario.mount != mount: + return RenderedResponse( + 400, + "application/json", + _json_bytes({"error": f"scenario {scenario_id} is wire {scenario.wire}, not mount {mount}"}), + ) + tail = "/".join(segments[2:]) + return _render(scenario, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + + +class _ScriptedHandler(BaseHTTPRequestHandler): + store: Final[_ScenarioStore] = _ScenarioStore() + + def _dispatch(self, method: str) -> None: + length = int(self.headers.get("content-length") or 0) + body = self.rfile.read(length) if length else b"" + rendered = handle_request(self.store, method, self.path, body) + self.send_response(rendered.status_code) + self.send_header("content-type", rendered.content_type) + self.send_header("content-length", str(len(rendered.body))) + self.end_headers() + self.wfile.write(rendered.body) + + def do_GET(self) -> None: + self._dispatch("GET") + + def do_POST(self) -> None: + self._dispatch("POST") + + def do_DELETE(self) -> None: + self._dispatch("DELETE") + + + +DEFAULT_PORT: Final = 9100 + + +def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None: + server = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) + sys.stderr.write(f"scripted-provider listening on http://{bind_host}:{port}\n") + server.serve_forever() + + +if __name__ == "__main__": + port_arg = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT + serve(port=port_arg) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py new file mode 100644 index 00000000000..8d7678cf9ca --- /dev/null +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -0,0 +1,115 @@ +"""Token-pricing e2e: every (frontier model, pricing-component case) cell runs a +scripted-usage call through a deployment registered on the cost-map proxy, and +the spend row plus response-cost header must equal literal arithmetic on the +test map's rates. + +Nothing here touches a real provider or the bundled cost map: the proxy's +upstream is the scripted-provider sidecar and its entire cost map is +tests/e2e/cost_map.json. +""" + +from __future__ import annotations + +import pytest + +from conftest import CostCalcClient, cost_rows, register_scenario_deployment +from cost_matrix import ( + FRONTIER_MODELS, + Case, + FrontierModel, + cases_for, + expected_cost, + expected_token_columns, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatStreamOptions + +pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] + +_MATRIX: list[tuple[FrontierModel, Case]] = [ + (model, case) for model in FRONTIER_MODELS for case in cases_for(model) +] + + +def _case_id(param: tuple[FrontierModel, Case]) -> str: + model, case = param + return f"{model.map_key.replace('/', '-')}-{case.name}" + + +def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: + return ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted pricing call")], + stream=case.stream, + stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, + service_tier=case.service_tier, + ) + + +class TestTokenPricing: + @pytest.mark.parametrize("model_case", _MATRIX, ids=_case_id) + @pytest.mark.covers("quota_management.spend_tracking.cost_matrix.logs_cost") + def test_scripted_usage_bills_at_map_rates( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + model_case: tuple[FrontierModel, Case], + ) -> None: + model, case = model_case + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=_chat_body(model_name, marker, case), + stream=case.stream, + ) + assert response.ok, ( + f"{model.map_key}/{case.name}: proxy returned {response.status_code}: {response.body[:400]}" + ) + assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" + + expected = expected_cost(model, case) + if case.exact_spend and not case.stream: + # Streamed responses commit headers before the bill is computed, so + # the x-litellm-response-cost header is asserted only on non-stream + # calls. + assert response.response_cost is not None and cost_rows.approx_equal( + response.response_cost, expected + ), ( + f"x-litellm-response-cost {response.response_cost} != expected {expected}" + ) + + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}" + + if not case.exact_spend and case.expect_zero_bill: + # The provider reported no usage and this wire has no proxy-side + # recount, so the bill is exactly zero. + assert row.spend is not None and row.spend == 0, f"no-usage stream billed {row.spend}: {row}" + return + if not case.exact_spend: + # stream_usage=absent: the provider reported no usage, so the row's + # token counts are the proxy's own recount; only assert a bill landed. + assert row.spend is not None and row.spend > 0, f"no-usage stream billed nothing: {row}" + return + + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( + f"{model.map_key}/{case.name}: spend {row.spend} != expected {expected} " + f"(breakdown {row.breakdown.model_dump()})" + ) + + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py new file mode 100644 index 00000000000..b1ef675d9ef --- /dev/null +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -0,0 +1,186 @@ +"""Wire-format e2e: one scripted upstream per provider wire, answering with a +usage payload where every token kind the wire can report is nonzero. The spend +row's gross input cost must equal fresh tokens at the input rate plus each cache +and audio component at its own rate -- proving the wire's usage shape landed the +cached tokens inside the total (OpenAI/Gemini) or as separate fields +(Anthropic), and that the biller subtracted them before billing fresh tokens. + +Also covers the Responses API wire (an openai/gpt-5.5-pro deployment bridged by +the proxy to POST /responses) and a streamed Anthropic-messages case. +""" + +from __future__ import annotations + +import pytest + +from conftest import CostCalcClient, cost_rows, register_scenario_deployment +from cost_matrix import ( + FRONTIER_MODELS, + Case, + FrontierModel, + expected_breakdown, + expected_token_columns, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatStreamOptions +from scripted_provider import ScriptedUsage + +pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] + +_MODELS: dict[str, FrontierModel] = {model.map_key: model for model in FRONTIER_MODELS} + +# One scripted usage per wire, every reportable token kind nonzero. +_WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { + "openai_chat": ( + "gpt-5.6", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "openai_responses": ( + "gpt-5.5-pro", + ScriptedUsage( + fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25, reasoning_tokens=15 + ), + ), + "anthropic_messages": ( + "claude-sonnet-5", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + ), + ), + "gemini_generate": ( + "gemini/gemini-3.8-flash", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "together_chat": ( + "together_ai/moonshotai/Kimi-K3", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "fireworks_chat": ( + "fireworks_ai/kimi-k3", + ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), + ), +} + + +class TestWireFormats: + @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_wire_usage_shape_bills_each_component( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + wire: str, + ) -> None: + map_key, usage = _WIRE_USAGE[wire] + model = _MODELS[map_key] + case = Case(name="basic", usage=usage) + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted wire call")], + ), + ) + assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" + + expected = expected_breakdown(model, case) + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"{wire}: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"{wire}: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + breakdown = row.breakdown + assert breakdown.input_cost is not None and cost_rows.approx_equal( + breakdown.input_cost, expected.input_cost + ), ( + f"{wire}: gross input_cost {breakdown.input_cost} != expected {expected.input_cost}; " + "cached/written tokens billed at the input rate" + ) + assert breakdown.output_cost is not None and cost_rows.approx_equal( + breakdown.output_cost, expected.output_cost + ), f"{wire}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" + + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"{wire}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"{wire}: completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) + + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_anthropic_streamed_usage_bills_each_component( + self, client: CostCalcClient, resources: ResourceManager, scoped_key: str + ) -> None: + map_key, usage = _WIRE_USAGE["anthropic_messages"] + model = _MODELS[map_key] + case = Case(name="stream", usage=usage, stream=True) + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted anthropic stream")], + stream=True, + stream_options=ChatStreamOptions(include_usage=True), + ), + stream=True, + ) + assert response.ok, f"anthropic stream: proxy returned {response.status_code}: {response.body[:400]}" + assert response.stream_done, "anthropic stream did not reach its terminal event" + assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" + + expected = expected_breakdown(model, case) + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, "anthropic stream: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"anthropic stream: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json new file mode 100644 index 00000000000..b761710bae3 --- /dev/null +++ b/tests/e2e/cost_map.json @@ -0,0 +1,352 @@ +{ + "claude-haiku-4-5": { + "cache_creation_input_token_cost": 0.00021, + "cache_creation_input_token_cost_above_1hr": 0.00028000000000000003, + "cache_read_input_token_cost": 7e-06, + "input_cost_per_token": 7.000000000000001e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00014000000000000001, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "claude-opus-5": { + "cache_creation_input_token_cost": 0.00015000000000000001, + "cache_creation_input_token_cost_above_1hr": 0.0002, + "cache_read_input_token_cost": 4.9999999999999996e-06, + "input_cost_per_token": 5e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "claude-sonnet-5": { + "cache_creation_input_token_cost": 0.00018, + "cache_creation_input_token_cost_above_1hr": 0.00024000000000000003, + "cache_read_input_token_cost": 6e-06, + "input_cost_per_token": 6.000000000000001e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00012000000000000002, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "fireworks_ai/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 1.4e-05, + "input_cost_per_token": 0.00014000000000000001, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00028000000000000003, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "fireworks_ai/kimi-k3": { + "cache_read_input_token_cost": 1.2e-05, + "input_cost_per_token": 0.00012000000000000002, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00024000000000000003, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "fireworks_ai/qwen3p8-max": { + "cache_read_input_token_cost": 1.3e-05, + "input_cost_per_token": 0.00013000000000000002, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00026000000000000003, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gemini/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 9e-06, + "input_cost_per_audio_token": 0.00054, + "input_cost_per_token": 9e-05, + "input_cost_per_token_above_200k_tokens": 0.00072, + "input_cost_per_token_flex": 0.000135, + "input_cost_per_token_priority": 0.000153, + "litellm_provider": "gemini", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0006299999999999999, + "output_cost_per_reasoning_token": 0.00045000000000000004, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_200k_tokens": 0.0008100000000000001, + "output_cost_per_token_flex": 0.00022500000000000002, + "output_cost_per_token_priority": 0.000243, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-3.8-flash": { + "cache_read_input_token_cost": 8e-06, + "input_cost_per_audio_token": 0.00048, + "input_cost_per_token": 8e-05, + "input_cost_per_token_above_200k_tokens": 0.00064, + "input_cost_per_token_flex": 0.00012, + "input_cost_per_token_priority": 0.000136, + "litellm_provider": "gemini", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00056, + "output_cost_per_reasoning_token": 0.0004, + "output_cost_per_token": 0.00016, + "output_cost_per_token_above_200k_tokens": 0.00072, + "output_cost_per_token_flex": 0.0002, + "output_cost_per_token_priority": 0.000216, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 3e-06, + "input_cost_per_token": 3.0000000000000004e-05, + "input_cost_per_token_above_200k_tokens": 0.00024000000000000003, + "input_cost_per_token_flex": 4.5e-05, + "input_cost_per_token_priority": 5.1e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_reasoning_token": 0.00015000000000000001, + "output_cost_per_token": 6.000000000000001e-05, + "output_cost_per_token_above_200k_tokens": 0.00027, + "output_cost_per_token_flex": 7.500000000000001e-05, + "output_cost_per_token_priority": 8.099999999999999e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00012, + "cache_creation_input_token_cost_above_1hr": 0.00016, + "cache_read_input_token_cost": 4e-06, + "input_cost_per_audio_token": 0.00024, + "input_cost_per_token": 4e-05, + "input_cost_per_token_above_200k_tokens": 0.00032, + "input_cost_per_token_flex": 6e-05, + "input_cost_per_token_priority": 6.8e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00028, + "output_cost_per_reasoning_token": 0.0002, + "output_cost_per_token": 8e-05, + "output_cost_per_token_above_200k_tokens": 0.00036, + "output_cost_per_token_flex": 0.0001, + "output_cost_per_token_priority": 0.000108, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "input_cost_per_token_above_200k_tokens": 0.00016, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_priority": 3.4e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_reasoning_token": 0.0001, + "output_cost_per_token": 4e-05, + "output_cost_per_token_above_200k_tokens": 0.00018, + "output_cost_per_token_flex": 5e-05, + "output_cost_per_token_priority": 5.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gpt-5.6": { + "cache_creation_input_token_cost": 3e-05, + "cache_creation_input_token_cost_above_1hr": 4e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_audio_token": 6e-05, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_200k_tokens": 8e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_priority": 1.7e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 7e-05, + "output_cost_per_reasoning_token": 5e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_200k_tokens": 9e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 2.7e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "together_ai/moonshotai/Kimi-K3": { + "cache_creation_input_token_cost": 0.00030000000000000003, + "cache_creation_input_token_cost_above_1hr": 0.0004, + "cache_read_input_token_cost": 9.999999999999999e-06, + "input_cost_per_audio_token": 0.0006000000000000001, + "input_cost_per_token": 0.0001, + "input_cost_per_token_above_200k_tokens": 0.0008, + "input_cost_per_token_flex": 0.00015000000000000001, + "input_cost_per_token_priority": 0.00017, + "litellm_provider": "together_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0006999999999999999, + "output_cost_per_reasoning_token": 0.0005, + "output_cost_per_token": 0.0002, + "output_cost_per_token_above_200k_tokens": 0.0009000000000000001, + "output_cost_per_token_flex": 0.00025, + "output_cost_per_token_priority": 0.00027, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "together_ai/zai-org/GLM-5.3": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, + "cache_read_input_token_cost": 1.1e-05, + "input_cost_per_audio_token": 0.00066, + "input_cost_per_token": 0.00011, + "input_cost_per_token_above_200k_tokens": 0.00088, + "input_cost_per_token_flex": 0.000165, + "input_cost_per_token_priority": 0.000187, + "litellm_provider": "together_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, + "output_cost_per_token": 0.00022, + "output_cost_per_token_above_200k_tokens": 0.00099, + "output_cost_per_token_flex": 0.000275, + "output_cost_per_token_priority": 0.000297, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + } +} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index ad0914d455b..6b40e70125c 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -63,3 +63,5 @@ - {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} - {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} - {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} +- {id: quota_management.spend_tracking.cost_matrix.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_matrix, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "A scripted-usage call through the cost-map proxy bills every reported token kind at the deployment's test-map rate (input, output, cache read, 5m/1h cache write, reasoning, audio, above-threshold tiers, flex/priority service tiers, web search, response-model override) and lands on the row's cost_breakdown, streamed or not"} +- {id: quota_management.spend_tracking.scripted_wire.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: scripted_wire, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "Each provider wire shape (openai chat, responses, anthropic messages, gemini generateContent, together, fireworks) parses usage into the same spend components: the gross input cost is fresh tokens at the input rate plus each cache/audio component at its own rate, streamed anthropic included"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..a891d9dcba2 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,6 +145,22 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +# The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL +# pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a +# scripted-provider sidecar; deselected unless the opt-in env var is set. +COST_MAP_OPT_IN_ENV = "E2E_COST_MAP_STACK" +# Base URL of the proxy running the test cost map. Defaults to the shared proxy +# so a local run only has to set the opt-in and boot the proxy accordingly. +COST_MAP_PROXY_URL = os.environ.get("E2E_COST_MAP_PROXY_URL", PROXY_BASE_URL).rstrip("/") +# Where the test runner reaches the scripted-provider sidecar's control API. +SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get( + "E2E_SCRIPTED_PROVIDER_CONTROL_URL", "http://127.0.0.1:9100" +).rstrip("/") +# The api_base root deployments register with: how the proxy (possibly in +# another container) reaches the sidecar's provider wire. +SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get( + "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL +).rstrip("/") ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..7d37bcc6d3e 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -11,3 +11,4 @@ markers = managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set + cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json (LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless E2E_COST_MAP_STACK is set From 269afbe382df06d33780571a40a55e527afea2b7 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:28:44 +0000 Subject: [PATCH 033/349] test(e2e): apply review nits to cost calculation suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 28 +- tests/e2e/cost_calculation/cost_matrix.py | 174 ++-- tests/e2e/cost_calculation/scripted_client.py | 12 +- .../e2e/cost_calculation/scripted_provider.py | 803 +++++++++++------- .../test_token_pricing_e2e.py | 17 +- .../cost_calculation/test_wire_formats_e2e.py | 43 +- 6 files changed, 620 insertions(+), 457 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 1bba3d50e1d..345ca26f7e3 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -13,7 +13,7 @@ from __future__ import annotations import importlib.util import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path from types import ModuleType @@ -34,16 +34,16 @@ def _load_cost_rows() -> ModuleType: """Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree has no package layout), the same trick the mcp suite uses for logging/datadog_reader.py.""" - path = ( + path: Final = ( Path(__file__).resolve().parent.parent / "quota_management" / "spend_tracking" / "cost_rows.py" ) - name = "e2e_spend_tracking_cost_rows" - spec = importlib.util.spec_from_file_location(name, path) + name: Final = "e2e_spend_tracking_cost_rows" + spec: Final = importlib.util.spec_from_file_location(name, path) assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) + module: Final = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) return module @@ -59,7 +59,7 @@ class SpendCostBreakdown(Protocol): total_cost: float | None service_tier: str | None - def model_dump(self) -> dict[str, object]: ... + def model_dump(self) -> Mapping[str, object]: ... class SpendRowMetadata(Protocol): @@ -89,7 +89,9 @@ class CostRowsModule(Protocol): ] -cost_rows: Final[CostRowsModule] = cast(CostRowsModule, _load_cost_rows()) +cost_rows: Final[CostRowsModule] = cast( # cast-ok: cost_rows.py is loaded by path, so basedpyright has no importable name for it; its surface is declared in CostRowsModule + CostRowsModule, _load_cost_rows() +) @dataclass(frozen=True, slots=True) @@ -101,7 +103,7 @@ class CostCalcClient: @pytest.fixture(scope="session") def client() -> CostCalcClient: - proxy = build_proxy_client( + proxy: Final = build_proxy_client( base_url=COST_MAP_PROXY_URL, control_plane_base_url=COST_MAP_PROXY_URL, replica_urls=(COST_MAP_PROXY_URL,), @@ -118,18 +120,18 @@ def register_scenario_deployment( ) -> tuple[str, ScenarioHandle]: """Register the case's scenario on the sidecar plus a deployment pointed at it; both are torn down by ``resources``. Returns the callable model_name.""" - scenario: Scenario = case.scenario( + scenario: Final[Scenario] = case.scenario( scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}" ) - handle = register_scenario(scenario) + handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) - model_name = f"{model.model_name}-{marker}" - model_id = client.proxy.register_model( + model_name: Final = f"{model.model_name}-{marker}" + model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, litellm_params=LiteLLMParamsBody( model=model.litellm_model, - api_key="sk-scripted-provider", + api_key=model.api_key, api_base=handle.api_base(), ), model_info=ModelInfoBody(), diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index bc466d7d823..e8b1d249559 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -18,9 +18,11 @@ creation), the case is absent from the matrix rather than silently zero. from __future__ import annotations import json +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -64,8 +66,8 @@ class CostMapEntry(BaseModel): _COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) -_COST_MAP: Final[dict[str, CostMapEntry]] = _COST_MAP_ADAPTER.validate_python( - json.loads(COST_MAP_PATH.read_text()) +_COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( + _COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) ) TIER_THRESHOLD_TOKENS: Final = 200_000 @@ -109,7 +111,7 @@ class FrontierModel: # Response-model override targets: emit a sibling's bare provider-facing name so # the biller's provider-prefixed lookup lands on that sibling's map key. -_OVERRIDE_MODELS: Final[dict[str, str]] = { +_OVERRIDE_MODELS: Final[Mapping[str, str]] = MappingProxyType({ "gpt-5.6": "gpt-5.4-mini", "gpt-5.5-pro": "gpt-5.3-codex", "gpt-5.3-codex": "gpt-5.5-pro", @@ -124,9 +126,9 @@ _OVERRIDE_MODELS: Final[dict[str, str]] = { "fireworks_ai/kimi-k3": "qwen3p8-max", "fireworks_ai/qwen3p8-max": "kimi-k3", "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", -} +}) -_OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { +_OVERRIDE_MAP_KEYS: Final[Mapping[str, str]] = MappingProxyType({ "gpt-5.4-mini": "gpt-5.4-mini", "gpt-5.6": "gpt-5.6", "gpt-5.3-codex": "gpt-5.3-codex", @@ -139,7 +141,7 @@ _OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", "qwen3p8-max": "fireworks_ai/qwen3p8-max", "kimi-k3": "fireworks_ai/kimi-k3", -} +}) _FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( @@ -176,7 +178,7 @@ def _frontier() -> tuple[FrontierModel, ...]: FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() # Token kinds each wire can report, gating which pricing cases apply. -_WIRE_CAPS: Final[dict[str, frozenset[str]]] = { +_WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "openai_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", @@ -205,9 +207,9 @@ _WIRE_CAPS: Final[dict[str, frozenset[str]]] = { "web_search", "response_model", "absent_usage", } ), -} +}) -CaseName = Literal[ +CaseName: TypeAlias = Literal[ "basic", "cache_read", "cache_write_5m", @@ -260,7 +262,7 @@ _BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) def _web_search_case(model: FrontierModel) -> Case: - counts_exactly = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") return Case( name="web_search", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), @@ -269,26 +271,24 @@ def _web_search_case(model: FrontierModel) -> Case: def cases_for(model: FrontierModel) -> tuple[Case, ...]: - rates = model.rates - caps = _WIRE_CAPS[model.wire] - cases: list[Case] = [Case(name="basic", usage=_BASIC_USAGE)] - if rates.cache_read_input_token_cost is not None and "cache_read" in caps: - cases.append( + rates: Final = model.rates + caps: Final = _WIRE_CAPS[model.wire] + candidates: Final[tuple[Case | None, ...]] = ( + Case(name="basic", usage=_BASIC_USAGE), + ( Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) - ) - if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps: - cases.append( + if rates.cache_read_input_token_cost is not None and "cache_read" in caps + else None + ), + ( Case( name="cache_write_5m", usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), ) - ) - if ( - rates.cache_creation_input_token_cost_above_1hr is not None - and rates.cache_creation_input_token_cost is not None - and "cache_write_1h" in caps - ): - cases.append( + if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps + else None + ), + ( Case( name="cache_write_1h", usage=ScriptedUsage( @@ -298,52 +298,61 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: output_tokens=30, ), ) - ) - if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps: - cases.append( + if ( + rates.cache_creation_input_token_cost_above_1hr is not None + and rates.cache_creation_input_token_cost is not None + and "cache_write_1h" in caps + ) + else None + ), + ( Case( name="reasoning", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), ) - ) - if ( - rates.input_cost_per_audio_token is not None - and rates.output_cost_per_audio_token is not None - and "audio" in caps - ): - cases.append( + if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps + else None + ), + ( Case( name="audio", usage=ScriptedUsage( fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 ), ) - ) - if ( - rates.input_cost_per_token_above_200k_tokens is not None - and rates.output_cost_per_token_above_200k_tokens is not None - ): - cases.append( + if ( + rates.input_cost_per_audio_token is not None + and rates.output_cost_per_audio_token is not None + and "audio" in caps + ) + else None + ), + ( Case( name="tiered", usage=ScriptedUsage( fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 ), ) - ) - if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None: - cases.append( + if ( + rates.input_cost_per_token_above_200k_tokens is not None + and rates.output_cost_per_token_above_200k_tokens is not None + ) + else None + ), + ( Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") - ) - if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None: - cases.append( + if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None + else None + ), + ( Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") - ) - if rates.search_context_cost_per_query is not None and "web_search" in caps: - cases.append(_web_search_case(model)) - cases.append(Case(name="stream", usage=_BASIC_USAGE, stream=True)) - if "absent_usage" in caps: - cases.append( + if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None + else None + ), + _web_search_case(model) if rates.search_context_cost_per_query is not None and "web_search" in caps else None, + Case(name="stream", usage=_BASIC_USAGE, stream=True), + ( Case( name="stream_no_usage", usage=_BASIC_USAGE, @@ -355,10 +364,16 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: # wires recount tokens proxy-side and bill a nonzero amount. expect_zero_bill=model.wire == "openai_responses", ) - ) - if "response_model" in caps: - cases.append(Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True)) - return tuple(cases) + if "absent_usage" in caps + else None + ), + ( + Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True) + if "response_model" in caps + else None + ), + ) + return tuple(case for case in candidates if case is not None) @dataclass(frozen=True, slots=True) @@ -387,38 +402,41 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: to the tier's variants, falling back to the base rate when a variant is unset -- mirroring _get_token_base_cost in litellm's cost calculator. """ - rates = model.override_rates if case.response_model_override else model.rates - u = case.usage - prompt_tokens = ( + rates: Final = model.override_rates if case.response_model_override else model.rates + u: Final = case.usage + prompt_tokens: Final = ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens + u.audio_input_tokens ) - tiered = prompt_tokens > TIER_THRESHOLD_TOKENS - in_rate = rates.input_cost_per_token or 0.0 - out_rate = rates.output_cost_per_token or 0.0 - if case.service_tier == "flex": - in_rate = rates.input_cost_per_token_flex or in_rate - out_rate = rates.output_cost_per_token_flex or out_rate - if case.service_tier == "priority": - in_rate = rates.input_cost_per_token_priority or in_rate - out_rate = rates.output_cost_per_token_priority or out_rate - if tiered: - in_rate = rates.input_cost_per_token_above_200k_tokens or in_rate - out_rate = rates.output_cost_per_token_above_200k_tokens or out_rate - input_cost = ( + tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate: Final = ( + (rates.input_cost_per_token_above_200k_tokens if tiered else None) + or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.input_cost_per_token + or 0.0 + ) + out_rate: Final = ( + (rates.output_cost_per_token_above_200k_tokens if tiered else None) + or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.output_cost_per_token + or 0.0 + ) + input_cost: Final = ( u.fresh_input_tokens * in_rate + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) - output_cost = ( + output_cost: Final = ( u.output_tokens * out_rate + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) ) - search = rates.search_context_cost_per_query - tool_cost = case.billed_web_search_calls * ( + search: Final = rates.search_context_cost_per_query + tool_cost: Final = case.billed_web_search_calls * ( search.search_context_size_medium if search and search.search_context_size_medium else 0.0 ) return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) @@ -432,7 +450,7 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: """(prompt_tokens, completion_tokens) the spend row should carry, per the wire's normalization: Anthropic folds cache read/write into prompt_tokens, everyone else reports the totals the wire emitted.""" - u = case.usage + u: Final = case.usage if model.wire == "anthropic_messages": return ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, diff --git a/tests/e2e/cost_calculation/scripted_client.py b/tests/e2e/cost_calculation/scripted_client.py index dceec02630a..9dbf9c98986 100644 --- a/tests/e2e/cost_calculation/scripted_client.py +++ b/tests/e2e/cost_calculation/scripted_client.py @@ -12,6 +12,7 @@ from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BA from e2e_http import URL, NoBody, unwrap, post from e2e_http import delete as http_delete from scripted_provider import ( + WIRE_MOUNTS, Scenario, ScenarioDeleted, ScenarioRegistered, @@ -29,19 +30,12 @@ class ScenarioHandle: return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}" def _mount(self) -> str: - return { - "openai_chat": "openai", - "openai_responses": "openai", - "anthropic_messages": "anthropic", - "gemini_generate": "gemini", - "together_chat": "together", - "fireworks_chat": "fireworks", - }[self.wire] + return WIRE_MOUNTS[self.wire] def register_scenario(scenario: Scenario) -> ScenarioHandle: """POST the scenario to the sidecar's control API and return its handle.""" - result = unwrap( + result: Final = unwrap( post( URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"), headers=NoBody(), diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 93a6f49ec25..f1deafd1bc5 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -31,14 +31,16 @@ import json import sys import threading import time +from collections.abc import Mapping from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError -Wire = Literal[ +Wire: TypeAlias = Literal[ "openai_chat", "openai_responses", "anthropic_messages", @@ -47,17 +49,19 @@ Wire = Literal[ "fireworks_chat", ] -_WIRE_MOUNTS: Final[dict[str, str]] = { - "openai_chat": "openai", - "openai_responses": "openai", - "anthropic_messages": "anthropic", - "gemini_generate": "gemini", - "together_chat": "together", - "fireworks_chat": "fireworks", -} +WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( + { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", + } +) -StreamUsage = Literal["final_chunk", "absent"] -ServiceTier = Literal["flex", "priority"] +StreamUsage: TypeAlias = Literal["final_chunk", "absent"] +ServiceTier: TypeAlias = Literal["flex", "priority"] class ScriptedUsage(BaseModel): @@ -106,7 +110,7 @@ class Scenario(BaseModel): @property def mount(self) -> str: - return _WIRE_MOUNTS[self.wire] + return WIRE_MOUNTS[self.wire] class ScenarioRegistered(BaseModel): @@ -128,369 +132,494 @@ class RenderedResponse: body: bytes -def _json_bytes(payload: dict[str, object]) -> bytes: - return json.dumps(payload).encode("utf-8") +def _jobj(*pairs: tuple[str, object]) -> Mapping[str, object]: + """A JSON object payload built in one shot and frozen.""" + return MappingProxyType(dict(pairs)) -def _sse(events: tuple[tuple[str | None, dict[str, object] | str], ...]) -> bytes: - frames: list[str] = [] - for event_name, data in events: - head = f"event: {event_name}\n" if event_name is not None else "" - payload = data if isinstance(data, str) else json.dumps(data) - frames.append(f"{head}data: {payload}\n\n") - return "".join(frames).encode("utf-8") +def _jobj_opt(*pairs: tuple[str, object] | None) -> Mapping[str, object]: + """``_jobj`` where a ``None`` pair means the field is absent.""" + return MappingProxyType(dict(pair for pair in pairs if pair is not None)) + + +def _json_bytes(payload: Mapping[str, object]) -> bytes: + return json.dumps(payload, default=dict).encode("utf-8") + + +def _sse_frame(event_name: str | None, data: Mapping[str, object] | str) -> str: + head: Final = f"event: {event_name}\n" if event_name is not None else "" + payload: Final = data if isinstance(data, str) else json.dumps(data, default=dict) + return f"{head}data: {payload}\n\n" + + +def _sse(events: tuple[tuple[str | None, Mapping[str, object] | str], ...]) -> bytes: + return "".join(_sse_frame(event_name, data) for event_name, data in events).encode("utf-8") # ---------- per-wire usage shapes ---------- -def _openai_usage(u: ScriptedUsage) -> dict[str, object]: - prompt_tokens = ( +def _openai_usage(u: ScriptedUsage) -> Mapping[str, object]: + prompt_tokens: Final = ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens + u.audio_input_tokens ) - completion_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - prompt_details: dict[str, object] = {} - if u.cache_read_tokens: - prompt_details["cached_tokens"] = u.cache_read_tokens - if u.cache_write_5m_tokens or u.cache_write_1h_tokens: - prompt_details["cache_write_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens - prompt_details["cache_creation_token_details"] = { - "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, - "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, - } - if u.audio_input_tokens: - prompt_details["audio_tokens"] = u.audio_input_tokens - completion_details: dict[str, object] = {} - if u.reasoning_tokens: - completion_details["reasoning_tokens"] = u.reasoning_tokens - if u.audio_output_tokens: - completion_details["audio_tokens"] = u.audio_output_tokens - usage: dict[str, object] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } - if prompt_details: - usage["prompt_tokens_details"] = prompt_details - if completion_details: - usage["completion_tokens_details"] = completion_details - return usage + completion_tokens: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + prompt_details: Final = _jobj_opt( + ("cached_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ( + ("cache_write_tokens", u.cache_write_5m_tokens + u.cache_write_1h_tokens) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ( + "cache_creation_token_details", + _jobj( + ("ephemeral_5m_input_tokens", u.cache_write_5m_tokens), + ("ephemeral_1h_input_tokens", u.cache_write_1h_tokens), + ), + ) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ("audio_tokens", u.audio_input_tokens) if u.audio_input_tokens else None, + ) + completion_details: Final = _jobj_opt( + ("reasoning_tokens", u.reasoning_tokens) if u.reasoning_tokens else None, + ("audio_tokens", u.audio_output_tokens) if u.audio_output_tokens else None, + ) + return _jobj_opt( + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("total_tokens", prompt_tokens + completion_tokens), + ("prompt_tokens_details", prompt_details) if prompt_details else None, + ("completion_tokens_details", completion_details) if completion_details else None, + ) -def _anthropic_usage(u: ScriptedUsage) -> dict[str, object]: +def _anthropic_usage(u: ScriptedUsage) -> Mapping[str, object]: # Anthropic reports uncached-only input_tokens; cache reads and writes ride # top-level fields, with the 5m/1h write split under cache_creation. - usage: dict[str, object] = { - "input_tokens": u.fresh_input_tokens, - "output_tokens": u.output_tokens, - } - if u.cache_read_tokens: - usage["cache_read_input_tokens"] = u.cache_read_tokens - if u.cache_write_5m_tokens or u.cache_write_1h_tokens: - usage["cache_creation_input_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens - usage["cache_creation"] = { - "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, - "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, - } - if u.web_search_calls: - usage["server_tool_use"] = {"web_search_requests": u.web_search_calls} - return usage + return _jobj_opt( + ("input_tokens", u.fresh_input_tokens), + ("output_tokens", u.output_tokens), + ("cache_read_input_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ( + ("cache_creation_input_tokens", u.cache_write_5m_tokens + u.cache_write_1h_tokens) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ( + "cache_creation", + _jobj( + ("ephemeral_5m_input_tokens", u.cache_write_5m_tokens), + ("ephemeral_1h_input_tokens", u.cache_write_1h_tokens), + ), + ) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ("server_tool_use", _jobj(("web_search_requests", u.web_search_calls))) + if u.web_search_calls + else None + ), + ) -def _gemini_usage(u: ScriptedUsage) -> dict[str, object]: +def _gemini_usage(u: ScriptedUsage) -> Mapping[str, object]: # promptTokenCount carries the cached count inside it; TEXT modality is the # cached-inclusive text count so litellm's implicit-caching subtraction lands # on the fresh figure. candidatesTokenCount includes reasoning + audio. - prompt_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens - candidates = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - usage: dict[str, object] = { - "promptTokenCount": prompt_tokens, - "candidatesTokenCount": candidates, - "totalTokenCount": prompt_tokens + candidates, - } - if u.cache_read_tokens: - usage["cachedContentTokenCount"] = u.cache_read_tokens - if u.reasoning_tokens: - usage["thoughtsTokenCount"] = u.reasoning_tokens - prompt_details = [{"modality": "TEXT", "tokenCount": u.fresh_input_tokens + u.cache_read_tokens}] - if u.audio_input_tokens: - prompt_details.append({"modality": "AUDIO", "tokenCount": u.audio_input_tokens}) - usage["promptTokensDetails"] = prompt_details - if u.audio_output_tokens: - usage["candidatesTokensDetails"] = [ - {"modality": "TEXT", "tokenCount": u.output_tokens + u.reasoning_tokens}, - {"modality": "AUDIO", "tokenCount": u.audio_output_tokens}, - ] - return usage + prompt_tokens: Final = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + candidates: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + return _jobj_opt( + ("promptTokenCount", prompt_tokens), + ("candidatesTokenCount", candidates), + ("totalTokenCount", prompt_tokens + candidates), + ("cachedContentTokenCount", u.cache_read_tokens) if u.cache_read_tokens else None, + ("thoughtsTokenCount", u.reasoning_tokens) if u.reasoning_tokens else None, + ( + "promptTokensDetails", + ( + _jobj(("modality", "TEXT"), ("tokenCount", u.fresh_input_tokens + u.cache_read_tokens)), + *( + (_jobj(("modality", "AUDIO"), ("tokenCount", u.audio_input_tokens)),) + if u.audio_input_tokens + else () + ), + ), + ), + ( + ( + "candidatesTokensDetails", + ( + _jobj(("modality", "TEXT"), ("tokenCount", u.output_tokens + u.reasoning_tokens)), + _jobj(("modality", "AUDIO"), ("tokenCount", u.audio_output_tokens)), + ), + ) + if u.audio_output_tokens + else None + ), + ) -def _responses_usage(u: ScriptedUsage) -> dict[str, object]: - input_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens - output_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - usage: dict[str, object] = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - input_details: dict[str, object] = {} - if u.cache_read_tokens: - input_details["cached_tokens"] = u.cache_read_tokens - if input_details: - usage["input_tokens_details"] = input_details - if u.reasoning_tokens: - usage["output_tokens_details"] = {"reasoning_tokens": u.reasoning_tokens} - return usage +def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: + input_tokens: Final = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + output_tokens: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + input_details: Final = _jobj_opt( + ("cached_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ) + return _jobj_opt( + ("input_tokens", input_tokens), + ("output_tokens", output_tokens), + ("total_tokens", input_tokens + output_tokens), + ("input_tokens_details", input_details) if input_details else None, + ( + ("output_tokens_details", _jobj(("reasoning_tokens", u.reasoning_tokens))) + if u.reasoning_tokens + else None + ), + ) # ---------- per-wire responses ---------- -def _openai_message(scenario: Scenario) -> dict[str, object]: - message: dict[str, object] = {"role": "assistant", "content": scenario.output.text} - if scenario.usage.web_search_calls: - message["annotations"] = [ - { - "type": "url_citation", - "url_citation": { - "url": "https://scripted.example/source", - "title": "scripted source", - "start_index": 0, - "end_index": 1, - }, - } - for _ in range(scenario.usage.web_search_calls) - ] - return message +def _openai_message(scenario: Scenario) -> Mapping[str, object]: + return _jobj_opt( + ("role", "assistant"), + ("content", scenario.output.text), + ( + ( + "annotations", + tuple( + _jobj( + ("type", "url_citation"), + ( + "url_citation", + _jobj( + ("url", "https://scripted.example/source"), + ("title", "scripted source"), + ("start_index", 0), + ("end_index", 1), + ), + ), + ) + for _ in range(scenario.usage.web_search_calls) + ), + ) + if scenario.usage.web_search_calls + else None + ), + ) -def _openai_chat_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - body: dict[str, object] = { - "id": f"chatcmpl-{scenario.scenario_id}", - "object": "chat.completion", - "created": int(time.time()), - "model": scenario.output.response_model or requested_model, - "choices": [ - { - "index": 0, - "message": _openai_message(scenario), - "finish_reason": scenario.output.finish_reason, - } - ], - "usage": _openai_usage(scenario.usage), - } - if scenario.service_tier is not None: - body["service_tier"] = scenario.service_tier - if scenario.output.provider_cost is not None: - body["cost"] = scenario.output.provider_cost - return body +def _openai_chat_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj_opt( + ("id", f"chatcmpl-{scenario.scenario_id}"), + ("object", "chat.completion"), + ("created", int(time.time())), + ("model", scenario.output.response_model or requested_model), + ( + "choices", + ( + _jobj( + ("index", 0), + ("message", _openai_message(scenario)), + ("finish_reason", scenario.output.finish_reason), + ), + ), + ), + ("usage", _openai_usage(scenario.usage)), + ("service_tier", scenario.service_tier) if scenario.service_tier is not None else None, + ("cost", scenario.output.provider_cost) if scenario.output.provider_cost is not None else None, + ) -def _openai_chunk(scenario: Scenario, requested_model: str, **kw: object) -> dict[str, object]: - chunk: dict[str, object] = { - "id": f"chatcmpl-{scenario.scenario_id}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": scenario.output.response_model or requested_model, - } - chunk.update(kw) - return chunk +def _openai_chunk( + scenario: Scenario, + requested_model: str, + choices: tuple[Mapping[str, object], ...] = (), + usage: Mapping[str, object] | None = None, +) -> Mapping[str, object]: + return _jobj_opt( + ("id", f"chatcmpl-{scenario.scenario_id}"), + ("object", "chat.completion.chunk"), + ("created", int(time.time())), + ("model", scenario.output.response_model or requested_model), + ("choices", choices), + ("usage", usage), + ) def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: - _EMPTY_DELTA: Final[dict[str, object]] = {} - delta: dict[str, object] = {"role": "assistant", "content": scenario.output.text} - if scenario.usage.web_search_calls: - delta["annotations"] = _openai_message(scenario)["annotations"] - events: list[tuple[str | None, dict[str, object] | str]] = [ + delta: Final = _jobj_opt( + ("role", "assistant"), + ("content", scenario.output.text), ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], - ), + ("annotations", _openai_message(scenario)["annotations"]) + if scenario.usage.web_search_calls + else None ), + ) + return _sse( ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", _jobj(("role", "assistant"))), ("finish_reason", None)),), + ), ), - ), - ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[ - { - "index": 0, - "delta": _EMPTY_DELTA, - "finish_reason": scenario.output.finish_reason, - } - ], + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", delta), ("finish_reason", None)),), + ), ), - ), - ] - if scenario.stream_usage == "final_chunk": - events.append( - (None, _openai_chunk(scenario, requested_model, choices=(), usage=_openai_usage(scenario.usage))) + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=( + _jobj( + ("index", 0), + ("delta", _jobj()), + ("finish_reason", scenario.output.finish_reason), + ), + ), + ), + ), + *( + ((None, _openai_chunk(scenario, requested_model, usage=_openai_usage(scenario.usage))),) + if scenario.stream_usage == "final_chunk" + else () + ), + (None, "[DONE]"), ) - events.append((None, "[DONE]")) - return _sse(tuple(events)) + ) -def _anthropic_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - return { - "id": f"msg_{scenario.scenario_id}", - "type": "message", - "role": "assistant", - "model": scenario.output.response_model or requested_model, - "content": [{"type": "text", "text": scenario.output.text}], - "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - "usage": _anthropic_usage(scenario.usage), - } +def _anthropic_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ("id", f"msg_{scenario.scenario_id}"), + ("type", "message"), + ("role", "assistant"), + ("model", scenario.output.response_model or requested_model), + ("content", (_jobj(("type", "text"), ("text", scenario.output.text)),)), + ( + "stop_reason", + "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + ), + ("usage", _anthropic_usage(scenario.usage)), + ) def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: - emit_usage = scenario.stream_usage == "final_chunk" - input_usage = {k: v for k, v in _anthropic_usage(scenario.usage).items() if k != "output_tokens"} - message_start: dict[str, object] = { - "type": "message_start", - "message": { - "id": f"msg_{scenario.scenario_id}", - "type": "message", - "role": "assistant", - "model": scenario.output.response_model or requested_model, - "content": [], - "stop_reason": None, - **({"usage": input_usage} if emit_usage else {}), - }, - } - message_delta: dict[str, object] = { - "type": "message_delta", - "delta": { - "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason - }, - **({"usage": {"output_tokens": scenario.usage.output_tokens}} if emit_usage else {}), - } + emit_usage: Final = scenario.stream_usage == "final_chunk" + input_usage: Final = _jobj( + *( + (key, value) + for key, value in _anthropic_usage(scenario.usage).items() + if key != "output_tokens" + ) + ) + message_start: Final = _jobj( + ("type", "message_start"), + ( + "message", + _jobj_opt( + ("id", f"msg_{scenario.scenario_id}"), + ("type", "message"), + ("role", "assistant"), + ("model", scenario.output.response_model or requested_model), + ("content", ()), + ("stop_reason", None), + ("usage", input_usage) if emit_usage else None, + ), + ), + ) + message_delta: Final = _jobj_opt( + ("type", "message_delta"), + ( + "delta", + _jobj( + ( + "stop_reason", + "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + ) + ), + ), + ( + ("usage", _jobj(("output_tokens", scenario.usage.output_tokens))) + if emit_usage + else None + ), + ) return _sse( ( ("message_start", message_start), ( "content_block_start", - { - "type": "content_block_start", - "index": 0, - "content_block": {"type": "text", "text": ""}, - }, + _jobj( + ("type", "content_block_start"), + ("index", 0), + ("content_block", _jobj(("type", "text"), ("text", ""))), + ), ), ( "content_block_delta", - { - "type": "content_block_delta", - "index": 0, - "delta": {"type": "text_delta", "text": scenario.output.text}, - }, + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), + ), ), - ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("content_block_stop", _jobj(("type", "content_block_stop"), ("index", 0))), ("message_delta", message_delta), - ("message_stop", {"type": "message_stop"}), + ("message_stop", _jobj(("type", "message_stop"))), ) ) -def _gemini_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - candidate: dict[str, object] = { - "content": {"parts": [{"text": scenario.output.text}], "role": "model"}, - "finishReason": "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), - "index": 0, - } - if scenario.usage.web_search_calls: - candidate["groundingMetadata"] = { - "webSearchQueries": [f"query {i}" for i in range(scenario.usage.web_search_calls)] - } - return { - "candidates": [candidate], - "usageMetadata": _gemini_usage(scenario.usage), - "modelVersion": scenario.output.response_model or requested_model, - } +def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ( + "candidates", + ( + _jobj_opt( + ( + "content", + _jobj( + ("parts", (_jobj(("text", scenario.output.text)),)), + ("role", "model"), + ), + ), + ( + "finishReason", + "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), + ), + ("index", 0), + ( + ( + "groundingMetadata", + _jobj( + ( + "webSearchQueries", + tuple(f"query {i}" for i in range(scenario.usage.web_search_calls)), + ) + ), + ) + if scenario.usage.web_search_calls + else None + ), + ), + ), + ), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ) def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: - first = _gemini_body(scenario, requested_model) - if scenario.stream_usage == "absent": - first = {k: v for k, v in first.items() if k != "usageMetadata"} - events: list[tuple[str | None, dict[str, object] | str]] = [(None, first)] - if scenario.stream_usage == "final_chunk": - events.append( - ( - None, - { - "candidates": [], - "usageMetadata": _gemini_usage(scenario.usage), - "modelVersion": scenario.output.response_model or requested_model, - }, - ) - ) - return _sse(tuple(events)) - - -def _responses_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - output: list[dict[str, object]] = [ - {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} - for i in range(scenario.usage.web_search_calls) - ] - output.append( - { - "type": "message", - "id": f"msg_{scenario.scenario_id}", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": scenario.output.text, - "annotations": [], - } - ], - } + emit_usage: Final = scenario.stream_usage == "final_chunk" + first: Final = ( + _jobj(*((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata")) + if scenario.stream_usage == "absent" + else _gemini_body(scenario, requested_model) + ) + return _sse( + ( + (None, first), + *( + ( + ( + None, + _jobj( + ("candidates", ()), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ), + ), + ) + if emit_usage + else () + ), + ) + ) + + +def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ("id", f"resp_{scenario.scenario_id}"), + ("object", "response"), + ("created_at", int(time.time())), + ("status", "completed"), + ("model", scenario.output.response_model or requested_model), + ( + "output", + ( + *( + _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) + for i in range(scenario.usage.web_search_calls) + ), + _jobj( + ("type", "message"), + ("id", f"msg_{scenario.scenario_id}"), + ("status", "completed"), + ("role", "assistant"), + ( + "content", + ( + _jobj( + ("type", "output_text"), + ("text", scenario.output.text), + ("annotations", ()), + ), + ), + ), + ), + ), + ), + ("usage", _responses_usage(scenario.usage)), ) - return { - "id": f"resp_{scenario.scenario_id}", - "object": "response", - "created_at": int(time.time()), - "status": "completed", - "model": scenario.output.response_model or requested_model, - "output": output, - "usage": _responses_usage(scenario.usage), - } def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: - completed = _responses_body(scenario, requested_model) - if scenario.stream_usage == "absent": - completed = {k: v for k, v in completed.items() if k != "usage"} - created = {**completed, "status": "in_progress", "usage": None} + completed: Final = ( + _jobj(*((key, value) for key, value in _responses_body(scenario, requested_model).items() if key != "usage")) + if scenario.stream_usage == "absent" + else _responses_body(scenario, requested_model) + ) + created: Final = _jobj( + *((key, value) for key, value in completed.items() if key not in ("status", "usage")), + ("status", "in_progress"), + ("usage", None), + ) return _sse( ( - ("response.created", {"type": "response.created", "response": created}), + ("response.created", _jobj(("type", "response.created"), ("response", created))), ( "response.output_text.delta", - { - "type": "response.output_text.delta", - "item_id": f"msg_{scenario.scenario_id}", - "output_index": scenario.usage.web_search_calls, - "content_index": 0, - "delta": scenario.output.text, - }, + _jobj( + ("type", "response.output_text.delta"), + ("item_id", f"msg_{scenario.scenario_id}"), + ("output_index", scenario.usage.web_search_calls), + ("content_index", 0), + ("delta", scenario.output.text), + ), ), - ("response.completed", {"type": "response.completed", "response": completed}), + ("response.completed", _jobj(("type", "response.completed"), ("response", completed))), ) ) @@ -538,11 +667,11 @@ class _ScenarioStore: _REQUEST_BODY: Final = TypeAdapter(dict[str, object]) -def _request_body(body: bytes) -> dict[str, object]: +def _request_body(body: bytes) -> Mapping[str, object]: try: return _REQUEST_BODY.validate_json(body) except ValueError: - return {} + return MappingProxyType({}) def _request_wants_stream(path_tail: str, body: bytes) -> bool: @@ -554,52 +683,66 @@ def _request_wants_stream(path_tail: str, body: bytes) -> bool: def _request_model(body: bytes) -> str: - model = _request_body(body).get("model") + model: Final = _request_body(body).get("model") return model if isinstance(model, str) else "unknown" def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: - path = urlsplit(raw_path).path - segments = [segment for segment in path.split("/") if segment] - if method == "GET" and segments == ["health"]: - return RenderedResponse(200, "application/json", _json_bytes({"status": "ok"})) + path: Final = urlsplit(raw_path).path + segments: Final = tuple(segment for segment in path.split("/") if segment) + if method == "GET" and segments == ("health",): + return RenderedResponse(200, "application/json", _json_bytes(_jobj(("status", "ok")))) if segments and segments[0] == "_scenarios": if method == "POST" and len(segments) == 1: try: - scenario = Scenario.model_validate_json(body) + scenario: Final = Scenario.model_validate_json(body) except ValidationError as exc: - return RenderedResponse(400, "application/json", _json_bytes({"error": str(exc)})) + return RenderedResponse( + 400, "application/json", _json_bytes(_jobj(("error", str(exc)))) + ) store.put(scenario) - return RenderedResponse(200, "application/json", _json_bytes({"scenario_id": scenario.scenario_id})) - if method == "DELETE" and len(segments) == 2: - deleted = store.drop(segments[1]) return RenderedResponse( - 200 if deleted else 404, "application/json", _json_bytes({"deleted": deleted}) + 200, "application/json", _json_bytes(_jobj(("scenario_id", scenario.scenario_id))) ) - return RenderedResponse(404, "application/json", _json_bytes({"error": "unknown control route"})) + if method == "DELETE" and len(segments) == 2: + deleted: Final = store.drop(segments[1]) + return RenderedResponse( + 200 if deleted else 404, + "application/json", + _json_bytes(_jobj(("deleted", deleted))), + ) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", "unknown control route"))) + ) if len(segments) < 2 or method != "POST": - return RenderedResponse(404, "application/json", _json_bytes({"error": f"no route for {method} {path}"})) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) + ) scenario_id, mount = segments[0], segments[1] - scenario = store.get(scenario_id) - if scenario is None: - return RenderedResponse(404, "application/json", _json_bytes({"error": f"unknown scenario {scenario_id}"})) - if scenario.mount != mount: + found: Final = store.get(scenario_id) + if found is None: + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", f"unknown scenario {scenario_id}"))) + ) + if found.mount != mount: return RenderedResponse( 400, "application/json", - _json_bytes({"error": f"scenario {scenario_id} is wire {scenario.wire}, not mount {mount}"}), + _json_bytes( + _jobj(("error", f"scenario {scenario_id} is wire {found.wire}, not mount {mount}")) + ), ) - tail = "/".join(segments[2:]) - return _render(scenario, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + tail: Final = "/".join(segments[2:]) + return _render(found, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) class _ScriptedHandler(BaseHTTPRequestHandler): store: Final[_ScenarioStore] = _ScenarioStore() def _dispatch(self, method: str) -> None: - length = int(self.headers.get("content-length") or 0) - body = self.rfile.read(length) if length else b"" - rendered = handle_request(self.store, method, self.path, body) + length: Final = int(self.headers.get("content-length") or 0) + body: Final = self.rfile.read(length) if length else b"" + rendered: Final = handle_request(self.store, method, self.path, body) self.send_response(rendered.status_code) self.send_header("content-type", rendered.content_type) self.send_header("content-length", str(len(rendered.body))) @@ -621,11 +764,11 @@ DEFAULT_PORT: Final = 9100 def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None: - server = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) + server: Final = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) sys.stderr.write(f"scripted-provider listening on http://{bind_host}:{port}\n") server.serve_forever() if __name__ == "__main__": - port_arg = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT + port_arg: Final = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT serve(port=port_arg) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 8d7678cf9ca..e210dad94b1 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -11,6 +11,7 @@ tests/e2e/cost_map.json. from __future__ import annotations import pytest +from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( @@ -25,11 +26,11 @@ from e2e_config import unique_marker from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatStreamOptions -pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] +pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark -_MATRIX: list[tuple[FrontierModel, Case]] = [ +_MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple( (model, case) for model in FRONTIER_MODELS for case in cases_for(model) -] +) def _case_id(param: tuple[FrontierModel, Case]) -> str: @@ -40,7 +41,7 @@ def _case_id(param: tuple[FrontierModel, Case]) -> str: def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: return ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted pricing call")], + messages=(ChatMessage(role="user", content=f"{marker} scripted pricing call"),), stream=case.stream, stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, service_tier=case.service_tier, @@ -58,9 +59,9 @@ class TestTokenPricing: model_case: tuple[FrontierModel, Case], ) -> None: model, case = model_case - marker = unique_marker() + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=_chat_body(model_name, marker, case), @@ -71,7 +72,7 @@ class TestTokenPricing: ) assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected = expected_cost(model, case) + expected: Final = expected_cost(model, case) if case.exact_spend and not case.stream: # Streamed responses commit headers before the bill is computed, so # the x-litellm-response-cost header is asserted only on non-stream @@ -82,7 +83,7 @@ class TestTokenPricing: f"x-litellm-response-cost {response.response_cost} != expected {expected}" ) - row = cost_rows.poll_cost_row_where( + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index b1ef675d9ef..c0276cf370c 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -12,6 +12,9 @@ the proxy to POST /responses) and a streamed Anthropic-messages case. from __future__ import annotations import pytest +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( @@ -26,12 +29,14 @@ from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatStreamOptions from scripted_provider import ScriptedUsage -pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] +pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark -_MODELS: dict[str, FrontierModel] = {model.map_key: model for model in FRONTIER_MODELS} +_MODELS: Final[Mapping[str, FrontierModel]] = MappingProxyType( + {model.map_key: model for model in FRONTIER_MODELS} +) # One scripted usage per wire, every reportable token kind nonzero. -_WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { +_WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ "openai_chat": ( "gpt-5.6", ScriptedUsage( @@ -89,7 +94,7 @@ _WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { "fireworks_ai/kimi-k3", ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), ), -} +}) class TestWireFormats: @@ -103,22 +108,22 @@ class TestWireFormats: wire: str, ) -> None: map_key, usage = _WIRE_USAGE[wire] - model = _MODELS[map_key] - case = Case(name="basic", usage=usage) - marker = unique_marker() + model: Final = _MODELS[map_key] + case: Final = Case(name="basic", usage=usage) + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted wire call")], + messages=(ChatMessage(role="user", content=f"{marker} scripted wire call"),), ), ) assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" - expected = expected_breakdown(model, case) - row = cost_rows.poll_cost_row_where( + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, @@ -128,7 +133,7 @@ class TestWireFormats: f"{wire}: spend {row.spend} != expected {expected.total} " f"(breakdown {row.breakdown.model_dump()})" ) - breakdown = row.breakdown + breakdown: Final = row.breakdown assert breakdown.input_cost is not None and cost_rows.approx_equal( breakdown.input_cost, expected.input_cost ), ( @@ -153,16 +158,16 @@ class TestWireFormats: self, client: CostCalcClient, resources: ResourceManager, scoped_key: str ) -> None: map_key, usage = _WIRE_USAGE["anthropic_messages"] - model = _MODELS[map_key] - case = Case(name="stream", usage=usage, stream=True) - marker = unique_marker() + model: Final = _MODELS[map_key] + case: Final = Case(name="stream", usage=usage, stream=True) + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted anthropic stream")], + messages=(ChatMessage(role="user", content=f"{marker} scripted anthropic stream"),), stream=True, stream_options=ChatStreamOptions(include_usage=True), ), @@ -172,8 +177,8 @@ class TestWireFormats: assert response.stream_done, "anthropic stream did not reach its terminal event" assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected = expected_breakdown(model, case) - row = cost_rows.poll_cost_row_where( + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, From 99bf8e9b2ffb6c647813029debba23c788e86b41 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:32:39 +0000 Subject: [PATCH 034/349] test(e2e): add cost calculation CI proxy config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/gateway/cost_calculation_ci_config.yml | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/gateway/cost_calculation_ci_config.yml diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index b6c3840f626..49cfc29aa17 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/gateway/cost_calculation_ci_config.yml b/tests/e2e/gateway/cost_calculation_ci_config.yml new file mode 100644 index 00000000000..ac0603fa7c1 --- /dev/null +++ b/tests/e2e/gateway/cost_calculation_ci_config.yml @@ -0,0 +1,7 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_batch_write_at: 5 + +model_list: [] From 415b06f5ff6d5959e2144ea826922694b2c8a60b Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 04:42:07 +0000 Subject: [PATCH 035/349] test(e2e): assert the real bill for the four fixed cost gaps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 22 +++++-------------- .../test_token_pricing_e2e.py | 5 ----- tests/e2e/cost_map.json | 15 +++++++++++++ 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index e8b1d249559..8f39e89a358 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -186,15 +186,12 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ } ), "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), - # Product gap: litellm hard-indexes message_delta["usage"] in - # anthropic/chat/handler.py, so a usage-absent anthropic stream raises - # KeyError; the real wire always carries it, so the case cannot be - # represented. - "anthropic_messages": frozenset({"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model"}), - # Product gap: the gemini transform sets ModelResponse.model from the - # request and drops the provider's modelVersion, so a response-model - # override can never be priced on this wire. - "gemini_generate": frozenset({"cache_read", "reasoning", "audio", "web_search", "absent_usage"}), + "anthropic_messages": frozenset( + {"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model", "absent_usage"} + ), + "gemini_generate": frozenset( + {"cache_read", "reasoning", "audio", "web_search", "response_model", "absent_usage"} + ), "together_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", @@ -240,9 +237,6 @@ class Case: billed_web_search_calls: int = 0 response_model_override: bool = False exact_spend: bool = True - # stream_usage=absent on a wire with no proxy-side token recount means the - # bill is exactly zero; asserted as such rather than skipped. - expect_zero_bill: bool = False def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( @@ -359,10 +353,6 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: stream=True, stream_usage="absent", exact_spend=False, - # The responses surface bills only provider-reported usage; - # with no usage in the stream the spend row is zero. Other - # wires recount tokens proxy-side and bill a nonzero amount. - expect_zero_bill=model.wire == "openai_responses", ) if "absent_usage" in caps else None diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index e210dad94b1..ead86931424 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -90,11 +90,6 @@ class TestTokenPricing: ) assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}" - if not case.exact_spend and case.expect_zero_bill: - # The provider reported no usage and this wire has no proxy-side - # recount, so the bill is exactly zero. - assert row.spend is not None and row.spend == 0, f"no-usage stream billed {row.spend}: {row}" - return if not case.exact_spend: # stream_usage=absent: the provider reported no usage, so the row's # token counts are the proxy's own recount; only assert a bill landed. diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index b761710bae3..68d840870c9 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -63,13 +63,18 @@ "supports_web_search": true }, "fireworks_ai/deepseek-v4p1-flash": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.4e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00014000000000000001, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00028000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -82,13 +87,18 @@ "supports_web_search": true }, "fireworks_ai/kimi-k3": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.2e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00012000000000000002, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00024000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -101,13 +111,18 @@ "supports_web_search": true }, "fireworks_ai/qwen3p8-max": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.3e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00013000000000000002, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00026000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, From 67778cfe2625eb97fd3d4733f1fae41aed7599fb Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:18:08 +0000 Subject: [PATCH 036/349] fix(cost): resolve dated openai/azure snapshots to their undated cost map entry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 9 +++++++- tests/test_litellm/test_cost_calculator.py | 24 ++++++++++++++++++++++ tests/test_litellm/test_utils.py | 19 +++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 734522c0c6a..33fdfcc36ba 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5286,6 +5286,13 @@ def _strip_stable_vertex_version(model_name) -> str: return re.sub(r"-\d+$", "", model_name) +_DATED_SNAPSHOT_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$") + + +def _strip_dated_snapshot_suffix(model_name: str) -> str: + return _DATED_SNAPSHOT_SUFFIX.sub("", model_name) + + def _get_base_bedrock_model(model_name) -> str: """ Get the base model from the given model name. @@ -5333,7 +5340,7 @@ def _strip_model_name(model: str, custom_llm_provider: str | None) -> str: strip_finetune: Final = _strip_openai_finetune_model_name(model_name=model) return strip_finetune else: - return model + return _strip_dated_snapshot_suffix(model_name=model) # Global case-insensitive lookup map for model_cost (built eagerly at module import) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7b53d3a58df..c5bc8d80fed 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -21,7 +21,9 @@ from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, + Choices, LiteLLMRealtimeStreamLoggingObject, + Message, ModelInfo, ModelResponse, PromptTokensDetailsWrapper, @@ -110,6 +112,28 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: + dated_response = ModelResponse( + model="gpt-5.6-luna-2026-07-09", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + dated_response._hidden_params = {"custom_llm_provider": "azure"} + + undated_response = ModelResponse( + model="gpt-5.6-luna", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + undated_response._hidden_params = {"custom_llm_provider": "azure"} + + dated_cost = litellm.completion_cost(completion_response=dated_response) + undated_cost = litellm.completion_cost(completion_response=undated_response) + + assert dated_cost == undated_cost + assert dated_cost > 0 + + def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): _hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 1000}} diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 46149589371..2f9e27af797 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -186,6 +186,25 @@ def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local assert info["key"] == "ft:gpt-4o-2024-08-06" +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_key"), + [ + ("gpt-5.6-luna-2026-07-09", "openai", "gpt-5.6-luna"), + ("gpt-5.6-luna-2026-07-09", "azure", "azure/gpt-5.6-luna"), + ], +) +def test_get_model_info_falls_back_from_dated_snapshot_to_undated_entry( + local_model_cost_map, model, custom_llm_provider, expected_key +): + info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + assert info["key"] == expected_key + + +def test_get_model_info_prefers_exact_dated_key_over_stripped(local_model_cost_map): + info = litellm.get_model_info(model="gpt-4o-2024-08-06", custom_llm_provider="openai") + assert info["key"] == "gpt-4o-2024-08-06" + + def test_check_provider_match_azure_ai_allows_openai_and_azure(): """ Test that azure_ai provider can match openai and azure models. From feb69c5f789d44a65dbbfa348ce39eaa3874b37f Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 16:38:15 +0000 Subject: [PATCH 037/349] test(e2e): add tool-call, terminal, and image-input shapes to the cost matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 184 +++++++- .../e2e/cost_calculation/scripted_provider.py | 408 +++++++++++++++--- .../test_token_pricing_e2e.py | 61 ++- .../cost_calculation/test_wire_formats_e2e.py | 111 ++++- 4 files changed, 688 insertions(+), 76 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 8f39e89a358..5f634712778 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -17,7 +17,11 @@ creation), the case is absent from the matrix rather than silently zero. from __future__ import annotations +import base64 import json +import random +import struct +import zlib from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -26,7 +30,7 @@ from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter -from scripted_provider import Scenario, ScriptedOutput, ScriptedUsage, Wire +from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" @@ -182,26 +186,37 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "openai_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", + } + ), + "openai_responses": frozenset( + { + "cache_read", "reasoning", "web_search", "response_model", "absent_usage", + "tool_call", "image_input", "responses_terminal", } ), - "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), "anthropic_messages": frozenset( - {"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model", "absent_usage"} + { + "cache_read", "cache_write_5m", "cache_write_1h", "web_search", + "response_model", "absent_usage", "tool_call", "image_input", + } ), "gemini_generate": frozenset( - {"cache_read", "reasoning", "audio", "web_search", "response_model", "absent_usage"} + { + "cache_read", "reasoning", "audio", "web_search", "response_model", + "absent_usage", "tool_call", "image_input", "prompt_blocked", + } ), "together_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), "fireworks_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), }) @@ -220,6 +235,15 @@ CaseName: TypeAlias = Literal[ "stream", "stream_no_usage", "response_model_override", + "stream_response_model_override", + "tool_call", + "stream_no_usage_tool_call", + "stream_no_usage_image_input", + "stream_no_usage_incomplete", + "stream_unvalidated", + "stream_no_usage_unvalidated", + "prompt_blocked", + "stream_prompt_blocked", ] @@ -237,6 +261,9 @@ class Case: billed_web_search_calls: int = 0 response_model_override: bool = False exact_spend: bool = True + tool_call: bool = False + image_input: bool = False + terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( @@ -246,6 +273,10 @@ class Case: output=ScriptedOutput( text=text, response_model=model.override_model if self.response_model_override else None, + tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) + if self.tool_call + else None, + terminal=self.terminal, ), stream_usage=self.stream_usage, service_tier=self.service_tier, @@ -254,6 +285,15 @@ class Case: _BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) +TOOL_CALL_ARGUMENTS: Final = json.dumps({ + "city": "Berlin", + "days": 7, + "units": "metric", + "notes": "filler " * 30, +}) + +_PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tokens=0) + def _web_search_case(model: FrontierModel) -> Case: counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") @@ -362,6 +402,100 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: if "response_model" in caps else None ), + ( + Case( + name="stream_response_model_override", + usage=_BASIC_USAGE, + stream=True, + response_model_override=True, + ) + if "response_model" in caps + else None + ), + ( + Case(name="tool_call", usage=_BASIC_USAGE, tool_call=True) + if "tool_call" in caps + else None + ), + ( + Case( + name="stream_no_usage_tool_call", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + tool_call=True, + exact_spend=False, + ) + if "absent_usage" in caps and "tool_call" in caps + else None + ), + ( + Case( + name="stream_no_usage_image_input", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + image_input=True, + exact_spend=False, + ) + if "absent_usage" in caps and "image_input" in caps + else None + ), + ( + Case( + name="stream_no_usage_incomplete", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + terminal="incomplete", + exact_spend=False, + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="stream_unvalidated", + usage=_BASIC_USAGE, + stream=True, + terminal="unvalidated", + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="stream_no_usage_unvalidated", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + terminal="unvalidated", + exact_spend=False, + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="prompt_blocked", + usage=_PROMPT_BLOCKED_USAGE, + terminal="prompt_blocked", + response_model_override=True, + ) + if "prompt_blocked" in caps + else None + ), + ( + Case( + name="stream_prompt_blocked", + usage=_PROMPT_BLOCKED_USAGE, + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ) + if "prompt_blocked" in caps + else None + ), ) return tuple(case for case in candidates if case is not None) @@ -436,6 +570,42 @@ def expected_cost(model: FrontierModel, case: Case) -> float: return expected_breakdown(model, case).total +def recount_cost( + model: FrontierModel, case: Case, prompt_tokens: int, completion_tokens: int +) -> float: + """What the proxy's own token recount should cost at the case's rates, + without pinning the tokenizer's exact counts.""" + rates: Final = model.override_rates if case.response_model_override else model.rates + return prompt_tokens * (rates.input_cost_per_token or 0.0) + completion_tokens * ( + rates.output_cost_per_token or 0.0 + ) + + +def _png_chunk(tag: bytes, payload: bytes) -> bytes: + return struct.pack(">I", len(payload)) + tag + payload + struct.pack(">I", zlib.crc32(tag + payload)) + + +def image_input_data_url() -> str: + """A deterministic 256x256 RGB noise PNG as a data URL; noise compresses + poorly on purpose so the base64 payload stays well above 100 KB and would + blow up the prompt recount if the URL were ever tokenized as text.""" + rng: Final = random.Random(0) + side: Final = 256 + raw: Final = b"".join( + b"\x00" + rng.randbytes(side * 3) for _ in range(side) + ) + png: Final = ( + b"\x89PNG\r\n\x1a\n" + + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", side, side, 8, 2, 0, 0, 0)) + + _png_chunk(b"IDAT", zlib.compress(raw)) + + _png_chunk(b"IEND", b"") + ) + return "data:image/png;base64," + base64.b64encode(png).decode() + + +IMAGE_INPUT_DATA_URL: Final = image_input_data_url() + + def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: """(prompt_tokens, completion_tokens) the spend row should carry, per the wire's normalization: Anthropic folds cache read/write into prompt_tokens, diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index f1deafd1bc5..e1a6c430307 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -38,7 +38,7 @@ from types import MappingProxyType from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator Wire: TypeAlias = Literal[ "openai_chat", @@ -62,6 +62,26 @@ WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( StreamUsage: TypeAlias = Literal["final_chunk", "absent"] ServiceTier: TypeAlias = Literal["flex", "priority"] +TerminalKind: TypeAlias = Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] + +# Which terminal variant each wire can represent. +_TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "openai_responses": frozenset({"incomplete", "unvalidated"}), + "gemini_generate": frozenset({"prompt_blocked"}), + } +) + + +class ScriptedToolCall(BaseModel): + """A single function call the scripted output emits instead of text. + ``arguments`` is the wire's JSON string (~250 chars), sliced into deltas + for streams.""" + + model_config = ConfigDict(frozen=True) + + name: str + arguments: str class ScriptedUsage(BaseModel): @@ -96,6 +116,12 @@ class ScriptedOutput(BaseModel): # OpenAI-compatible providers can report a provider-computed cost; emitted as # the top-level "cost" field on the together/fireworks wire. provider_cost: float | None = None + # When set, the response is a tool call only: no text content on any wire. + tool_call: ScriptedToolCall | None = None + # Terminal shape: "unvalidated" makes the Responses terminal response fail + # pydantic validation so the proxy takes its model_construct dict path; + # "prompt_blocked" is a Gemini promptFeedback-only body. + terminal: TerminalKind = "completed" class Scenario(BaseModel): @@ -108,6 +134,17 @@ class Scenario(BaseModel): stream_usage: StreamUsage = "final_chunk" service_tier: ServiceTier | None = None + @model_validator(mode="after") + def _check_terminal_supported(self) -> Scenario: + if ( + self.output.terminal != "completed" + and self.output.terminal not in _TERMINAL_CAPS.get(self.wire, frozenset()) + ): + raise ValueError( + f"wire {self.wire} cannot emit terminal={self.output.terminal}" + ) + return self + @property def mount(self) -> str: return WIRE_MOUNTS[self.wire] @@ -291,10 +328,38 @@ def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: # ---------- per-wire responses ---------- +def _split_arguments(arguments: str) -> tuple[str, ...]: + """Slice a tool-call arguments JSON string into 2-3 streamed deltas.""" + third: Final = max(1, len(arguments) // 3) + return tuple( + slice_ + for slice_ in (arguments[:third], arguments[third : 2 * third], arguments[2 * third :]) + if slice_ + ) + + def _openai_message(scenario: Scenario) -> Mapping[str, object]: + tool_call: Final = scenario.output.tool_call return _jobj_opt( ("role", "assistant"), - ("content", scenario.output.text), + ("content", None if tool_call is not None else scenario.output.text), + ( + ( + "tool_calls", + ( + _jobj( + ("id", f"call_{scenario.scenario_id}"), + ("type", "function"), + ( + "function", + _jobj(("name", tool_call.name), ("arguments", tool_call.arguments)), + ), + ), + ), + ) + if tool_call is not None + else None + ), ( ( "annotations", @@ -332,7 +397,12 @@ def _openai_chat_body(scenario: Scenario, requested_model: str) -> Mapping[str, _jobj( ("index", 0), ("message", _openai_message(scenario)), - ("finish_reason", scenario.output.finish_reason), + ( + "finish_reason", + "tool_calls" + if scenario.output.tool_call is not None + else scenario.output.finish_reason, + ), ), ), ), @@ -359,6 +429,7 @@ def _openai_chunk( def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: + tool_call: Final = scenario.output.tool_call delta: Final = _jobj_opt( ("role", "assistant"), ("content", scenario.output.text), @@ -368,6 +439,43 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: else None ), ) + body_deltas: Final[tuple[Mapping[str, object], ...]] = ( + ( + _jobj( + ("role", "assistant"), + ( + "tool_calls", + ( + _jobj( + ("index", 0), + ("id", f"call_{scenario.scenario_id}"), + ("type", "function"), + ( + "function", + _jobj(("name", tool_call.name), ("arguments", "")), + ), + ), + ), + ), + ), + *( + _jobj( + ( + "tool_calls", + ( + _jobj( + ("index", 0), + ("function", _jobj(("arguments", arguments_slice))), + ), + ), + ) + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ), + ) + if tool_call is not None + else (delta,) + ) return _sse( ( ( @@ -378,13 +486,16 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: choices=(_jobj(("index", 0), ("delta", _jobj(("role", "assistant"))), ("finish_reason", None)),), ), ), - ( - None, - _openai_chunk( - scenario, - requested_model, - choices=(_jobj(("index", 0), ("delta", delta), ("finish_reason", None)),), - ), + *( + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", body_delta), ("finish_reason", None)),), + ), + ) + for body_delta in body_deltas ), ( None, @@ -395,7 +506,12 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: _jobj( ("index", 0), ("delta", _jobj()), - ("finish_reason", scenario.output.finish_reason), + ( + "finish_reason", + "tool_calls" + if tool_call is not None + else scenario.output.finish_reason, + ), ), ), ), @@ -410,17 +526,34 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: ) +def _anthropic_content(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ("type", "tool_use"), + ("id", f"toolu_{scenario.scenario_id}"), + ("name", tool_call.name), + ("input", json.loads(tool_call.arguments)), + ), + ) + return (_jobj(("type", "text"), ("text", scenario.output.text)),) + + +def _anthropic_stop_reason(scenario: Scenario) -> str: + if scenario.output.tool_call is not None: + return "tool_use" + return "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + + def _anthropic_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: return _jobj( ("id", f"msg_{scenario.scenario_id}"), ("type", "message"), ("role", "assistant"), ("model", scenario.output.response_model or requested_model), - ("content", (_jobj(("type", "text"), ("text", scenario.output.text)),)), - ( - "stop_reason", - "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - ), + ("content", _anthropic_content(scenario)), + ("stop_reason", _anthropic_stop_reason(scenario)), ("usage", _anthropic_usage(scenario.usage)), ) @@ -453,12 +586,7 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: ("type", "message_delta"), ( "delta", - _jobj( - ( - "stop_reason", - "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - ) - ), + _jobj(("stop_reason", _anthropic_stop_reason(scenario))), ), ( ("usage", _jobj(("output_tokens", scenario.usage.output_tokens))) @@ -474,16 +602,45 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: _jobj( ("type", "content_block_start"), ("index", 0), - ("content_block", _jobj(("type", "text"), ("text", ""))), + ( + "content_block", + _jobj( + ("type", "tool_use"), + ("id", f"toolu_{scenario.scenario_id}"), + ("name", scenario.output.tool_call.name), + ("input", _jobj()), + ) + if scenario.output.tool_call is not None + else _jobj(("type", "text"), ("text", "")), + ), ), ), - ( - "content_block_delta", - _jobj( - ("type", "content_block_delta"), - ("index", 0), - ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), - ), + *( + tuple( + ( + "content_block_delta", + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ( + "delta", + _jobj(("type", "input_json_delta"), ("partial_json", arguments_slice)), + ), + ), + ) + for arguments_slice in _split_arguments(scenario.output.tool_call.arguments) + ) + if scenario.output.tool_call is not None + else ( + ( + "content_block_delta", + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), + ), + ), + ) ), ("content_block_stop", _jobj(("type", "content_block_stop"), ("index", 0))), ("message_delta", message_delta), @@ -492,7 +649,49 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: ) +def _gemini_prompt_blocked_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ( + "promptFeedback", + _jobj( + ("blockReason", "SAFETY"), + ( + "safetyRatings", + ( + _jobj( + ("category", "HARM_CATEGORY_HARASSMENT"), + ("probability", "HIGH"), + ("blocked", True), + ), + ), + ), + ), + ), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ) + + +def _gemini_parts(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ( + "functionCall", + _jobj( + ("name", tool_call.name), + ("args", json.loads(tool_call.arguments)), + ), + ) + ), + ) + return (_jobj(("text", scenario.output.text)),) + + def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + if scenario.output.terminal == "prompt_blocked": + return _gemini_prompt_blocked_body(scenario, requested_model) return _jobj( ( "candidates", @@ -501,7 +700,7 @@ def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, objec ( "content", _jobj( - ("parts", (_jobj(("text", scenario.output.text)),)), + ("parts", _gemini_parts(scenario)), ("role", "model"), ), ), @@ -559,67 +758,148 @@ def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: ) -def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: - return _jobj( - ("id", f"resp_{scenario.scenario_id}"), - ("object", "response"), - ("created_at", int(time.time())), - ("status", "completed"), - ("model", scenario.output.response_model or requested_model), - ( - "output", +def _responses_output(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + return ( + *( ( - *( - _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) - for i in range(scenario.usage.web_search_calls) - ), - _jobj( - ("type", "message"), - ("id", f"msg_{scenario.scenario_id}"), - ("status", "completed"), - ("role", "assistant"), - ( - "content", - ( - _jobj( - ("type", "output_text"), - ("text", scenario.output.text), - ("annotations", ()), - ), - ), + _jobj(("type", "scripted_future_item"), ("id", f"fut_{scenario.scenario_id}"), ("status", "completed")), + ) + if scenario.output.terminal == "unvalidated" + else () + ), + *( + _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) + for i in range(scenario.usage.web_search_calls) + ), + _jobj( + ("type", "function_call"), + ("id", f"fc_{scenario.scenario_id}"), + ("call_id", f"call_{scenario.scenario_id}"), + ("name", tool_call.name), + ("arguments", tool_call.arguments), + ("status", "completed"), + ) + if tool_call is not None + else _jobj( + ("type", "message"), + ("id", f"msg_{scenario.scenario_id}"), + ("status", "completed"), + ("role", "assistant"), + ( + "content", + ( + _jobj( + ("type", "output_text"), + ("text", scenario.output.text), + ("annotations", ()), ), ), ), ), + ) + + +def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + incomplete: Final = scenario.output.terminal == "incomplete" + return _jobj_opt( + ("id", f"resp_{scenario.scenario_id}"), + ("object", "response"), + ( + "created_at", + "not-a-number" if scenario.output.terminal == "unvalidated" else int(time.time()), + ), + ("status", "incomplete" if incomplete else "completed"), + ( + ("incomplete_details", _jobj(("reason", "max_output_tokens"))) + if incomplete + else None + ), + ("model", scenario.output.response_model or requested_model), + ("output", _responses_output(scenario)), ("usage", _responses_usage(scenario.usage)), ) def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: - completed: Final = ( + tool_call: Final = scenario.output.tool_call + terminal: Final = ( _jobj(*((key, value) for key, value in _responses_body(scenario, requested_model).items() if key != "usage")) if scenario.stream_usage == "absent" else _responses_body(scenario, requested_model) ) created: Final = _jobj( - *((key, value) for key, value in completed.items() if key not in ("status", "usage")), + *((key, value) for key, value in terminal.items() if key not in ("status", "usage")), ("status", "in_progress"), ("usage", None), ) - return _sse( + terminal_event: Final = ( + "response.incomplete" if scenario.output.terminal == "incomplete" else "response.completed" + ) + output_index: Final = ( + scenario.usage.web_search_calls + (1 if scenario.output.terminal == "unvalidated" else 0) + ) + middle_events: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( ( - ("response.created", _jobj(("type", "response.created"), ("response", created))), + ( + "response.output_item.added", + _jobj( + ("type", "response.output_item.added"), + ("output_index", output_index), + ( + "item", + _jobj( + ("type", "function_call"), + ("id", f"fc_{scenario.scenario_id}"), + ("call_id", f"call_{scenario.scenario_id}"), + ("name", tool_call.name), + ("arguments", ""), + ("status", "in_progress"), + ), + ), + ), + ), + *( + ( + "response.function_call_arguments.delta", + _jobj( + ("type", "response.function_call_arguments.delta"), + ("item_id", f"fc_{scenario.scenario_id}"), + ("output_index", output_index), + ("delta", arguments_slice), + ), + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ), + ( + "response.function_call_arguments.done", + _jobj( + ("type", "response.function_call_arguments.done"), + ("item_id", f"fc_{scenario.scenario_id}"), + ("output_index", output_index), + ("arguments", tool_call.arguments), + ), + ), + ) + if tool_call is not None + else ( ( "response.output_text.delta", _jobj( ("type", "response.output_text.delta"), ("item_id", f"msg_{scenario.scenario_id}"), - ("output_index", scenario.usage.web_search_calls), + ("output_index", output_index), ("content_index", 0), ("delta", scenario.output.text), ), ), - ("response.completed", _jobj(("type", "response.completed"), ("response", completed))), + ) + ) + return _sse( + ( + ("response.created", _jobj(("type", "response.created"), ("response", created))), + *middle_events, + (terminal_event, _jobj(("type", terminal_event), ("response", terminal))), ) ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index ead86931424..0b4f3e1fd37 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -16,15 +16,26 @@ from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( FRONTIER_MODELS, + IMAGE_INPUT_DATA_URL, Case, FrontierModel, cases_for, expected_cost, expected_token_columns, + recount_cost, ) from e2e_config import unique_marker from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions +from models import ( + ChatBody, + ChatMessage, + ChatStreamOptions, + ChatTool, + ChatToolFunction, + ImageContentPart, + ImageUrl, + TextContentPart, +) pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark @@ -41,10 +52,37 @@ def _case_id(param: tuple[FrontierModel, Case]) -> str: def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: return ChatBody( model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted pricing call"),), + messages=( + ChatMessage( + role="user", + content=( + [ + TextContentPart(text=f"{marker} scripted pricing call"), + ImageContentPart(image_url=ImageUrl(url=IMAGE_INPUT_DATA_URL)), + ] + if case.image_input + else f"{marker} scripted pricing call" + ), + ), + ), stream=case.stream, stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, service_tier=case.service_tier, + tools=( + ( + ChatTool( + function=ChatToolFunction( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ) + ), + ) + if case.tool_call + else None + ), ) @@ -92,8 +130,23 @@ class TestTokenPricing: if not case.exact_spend: # stream_usage=absent: the provider reported no usage, so the row's - # token counts are the proxy's own recount; only assert a bill landed. - assert row.spend is not None and row.spend > 0, f"no-usage stream billed nothing: {row}" + # token counts are the proxy's own recount; assert the recount + # billed both directions at the case's rates. + assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( + f"no-usage stream counted no input tokens: {row}" + ) + assert row.completion_tokens is not None and row.completion_tokens > 0, ( + f"no-usage stream counted no output tokens: {row}" + ) + if case.image_input: + assert row.prompt_tokens < 4000, ( + f"image data URL looks tokenized as text: prompt_tokens={row.prompt_tokens}" + ) + assert row.spend is not None and cost_rows.approx_equal( + row.spend, + recount_cost(model, case, row.prompt_tokens, row.completion_tokens), + ), f"no-usage stream spend {row.spend} != recount at map rates: {row}" + cost_rows.assert_total_is_sum_of_components(row) return assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index c0276cf370c..3c6c34b24fb 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -26,7 +26,7 @@ from cost_matrix import ( ) from e2e_config import unique_marker from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions +from models import ChatBody, ChatMessage, ChatStreamOptions, ChatTool, ChatToolFunction from scripted_provider import ScriptedUsage pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark @@ -96,6 +96,53 @@ _WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ ), }) +_SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) + +# Renderer-level shapes the pricing matrix gates per cap, pinned here once per +# wire so the sidecar emits prove they survive the proxy end to end. +_SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( + *( + ( + f"tool_call_{'stream' if stream else 'sync'}", + wire, + Case(name="tool_call", usage=_SHAPE_USAGE, stream=stream, tool_call=True), + ) + for wire in _WIRE_USAGE + for stream in (False, True) + ), + ( + "responses_incomplete", + "openai_responses", + Case(name="stream_no_usage_incomplete", usage=_SHAPE_USAGE, stream=True, terminal="incomplete"), + ), + ( + "responses_unvalidated", + "openai_responses", + Case(name="stream_unvalidated", usage=_SHAPE_USAGE, stream=True, terminal="unvalidated"), + ), + ( + "gemini_prompt_blocked", + "gemini_generate", + Case( + name="prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "gemini_prompt_blocked_stream", + "gemini_generate", + Case( + name="stream_prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ), + ), +) + class TestWireFormats: @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) @@ -189,3 +236,65 @@ class TestWireFormats: f"(breakdown {row.breakdown.model_dump()})" ) cost_rows.assert_total_is_sum_of_components(row) + + @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=lambda entry: entry[0]) + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_response_shape_bills_reported_usage( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + shape_wire_case: tuple[str, str, Case], + ) -> None: + shape, wire, case = shape_wire_case + map_key, _usage = _WIRE_USAGE[wire] + model: Final = _MODELS[map_key] + marker: Final = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response: Final = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=(ChatMessage(role="user", content=f"{marker} scripted {shape}"),), + stream=case.stream, + stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, + tools=( + ( + ChatTool( + function=ChatToolFunction( + name="get_weather", + parameters={"type": "object", "properties": {"city": {"type": "string"}}}, + ) + ), + ) + if case.tool_call + else None + ), + ), + stream=case.stream, + ) + assert response.ok, f"{shape}: proxy returned {response.status_code}: {response.body[:400]}" + if case.stream: + assert response.stream_done, f"{shape}: stream did not reach its terminal event" + assert response.stream_error is None, f"{shape}: stream error: {response.stream_error}" + + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"{shape}: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"{shape}: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"{shape}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"{shape}: completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) From 5507de326e3e98f9069af5f9d1315c89bb3c3e25 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 16:45:09 +0000 Subject: [PATCH 038/349] test(e2e): type the wire-shape parametrize ids callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/test_wire_formats_e2e.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index 3c6c34b24fb..4da7b31a6ef 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -144,6 +144,10 @@ _SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( ) +def _shape_id(entry: tuple[str, str, Case]) -> str: + return entry[0] + + class TestWireFormats: @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") @@ -237,7 +241,7 @@ class TestWireFormats: ) cost_rows.assert_total_is_sum_of_components(row) - @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=lambda entry: entry[0]) + @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=_shape_id) @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") def test_response_shape_bills_reported_usage( self, From 9885dc89621697e235fc65e0e86e147da87398c1 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:52:42 +0000 Subject: [PATCH 039/349] test(e2e): add azure, bedrock converse and vertex wires to the cost suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 54 +++- tests/e2e/cost_calculation/cost_matrix.py | 144 ++++++++- .../e2e/cost_calculation/scripted_provider.py | 284 +++++++++++++++++- .../cost_calculation/test_wire_formats_e2e.py | 64 ++++ tests/e2e/cost_map.json | 158 ++++++++++ tests/e2e/models.py | 1 + 6 files changed, 681 insertions(+), 24 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 345ca26f7e3..8c6db7c0010 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -12,6 +12,7 @@ Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). from __future__ import annotations import importlib.util +import json import sys from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -22,7 +23,7 @@ from typing import Final, Protocol, cast import pytest from cost_matrix import Case, FrontierModel -from e2e_config import COST_MAP_PROXY_URL +from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE from lifecycle import ResourceManager from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody from proxy_client import ProxyClient, build_proxy_client @@ -111,6 +112,41 @@ def client() -> CostCalcClient: return CostCalcClient(proxy=proxy) +_vertex_key_pem: str | None = None + + +def _vertex_service_account_json() -> str: + """A service-account credential JSON whose token_uri is the sidecar's + /_oauth/token route: the proxy's google-auth refresh then gets a scripted + access token without touching Google. One generated RSA key per process.""" + global _vertex_key_pem # mutable-ok: session-scoped key generation cached for reuse + if _vertex_key_pem is None: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + _vertex_key_pem = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": "cc-scripted-project", + "private_key_id": "scripted", + "private_key": _vertex_key_pem, + "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize", + "token_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/token", + } + ) + + def register_scenario_deployment( client: CostCalcClient, resources: ResourceManager, @@ -126,15 +162,21 @@ def register_scenario_deployment( handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) model_name: Final = f"{model.model_name}-{marker}" + extra_params: Final[dict[str, str]] = dict(model.litellm_params) + if model.wire == "vertex_generate": + extra_params["vertex_credentials"] = _vertex_service_account_json() model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, - litellm_params=LiteLLMParamsBody( - model=model.litellm_model, - api_key=model.api_key, - api_base=handle.api_base(), + litellm_params=LiteLLMParamsBody.model_validate( + { + "model": model.litellm_model, + "api_key": model.api_key, + "api_base": handle.api_base(), + **extra_params, + } ), - model_info=ModelInfoBody(), + model_info=ModelInfoBody(base_model=model.base_model), ) ) resources.defer(lambda: client.proxy.delete_model(model_id)) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 5f634712778..37495f37b0b 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -88,7 +88,14 @@ class FrontierModel: litellm_model: str wire: Wire map_key: str - override_model: str + override_model: str | None = None + override_map_key: str | None = None + # Registered as model_info.base_model; when set, the provider-reported + # model loses to it and every case bills at this deployment's own rates. + base_model: str | None = None + # Extra litellm_params merged into the /model/new registration (api_version, + # aws_* credentials, vertex_* auth). + litellm_params: Mapping[str, str] = MappingProxyType({}) @property def rates(self) -> CostMapEntry: @@ -96,11 +103,16 @@ class FrontierModel: @property def override_rates(self) -> CostMapEntry: + if self.base_model is not None or self.override_map_key is None: + return self.rates return _COST_MAP[self.override_map_key] @property - def override_map_key(self) -> str: - return _OVERRIDE_MAP_KEYS[self.override_model] + def provider_model(self) -> str: + """The bare provider-facing model name: litellm_model minus the provider + prefix and any routing segment (converse/, responses/).""" + tail: Final = self.litellm_model.split("/")[1:] + return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) @property def provider(self) -> str: @@ -166,6 +178,92 @@ _FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( ) +@dataclass(frozen=True, slots=True) +class _ExtendedSpec: + """A frontier entry whose override target, model_info.base_model or extra + litellm_params can't be derived from the map key alone.""" + + map_key: str + litellm_model: str + wire: Wire + override_model: str | None = None + override_map_key: str | None = None + base_model: str | None = None + litellm_params: Mapping[str, str] = MappingProxyType({}) + + +_AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"}) +_BEDROCK_PARAMS: Final[Mapping[str, str]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType( + { + "vertex_project": "cc-scripted-project", + "vertex_location": "us-central1", + } +) + +_EXTENDED_SPECS: Final[tuple[_ExtendedSpec, ...]] = ( + _ExtendedSpec( + map_key="azure/gpt-5.6", + litellm_model="azure/gpt-5.6", + wire="azure_chat", + override_model="gpt-5.4-mini", + override_map_key="azure/gpt-5.4-mini", + litellm_params=_AZURE_PARAMS, + ), + _ExtendedSpec( + # Deployment name is not a model; base_model pins billing so the + # response's model field loses, proving base_model wins. + map_key="azure/gpt-5.4-mini", + litellm_model="azure/cc-pinned-deployment", + wire="azure_chat", + override_model="gpt-5.6", + override_map_key="azure/gpt-5.6", + base_model="azure/gpt-5.4-mini", + litellm_params=_AZURE_PARAMS, + ), + _ExtendedSpec( + map_key="anthropic.claude-sonnet-5-v1:0", + litellm_model="bedrock/converse/anthropic.claude-sonnet-5-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="us.anthropic.claude-opus-5-v1:0", + litellm_model="bedrock/converse/us.anthropic.claude-opus-5-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="meta.llama4-maverick-17b-instruct-v1:0", + litellm_model="bedrock/converse/meta.llama4-maverick-17b-instruct-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="gemini-3.8-flash", + litellm_model="vertex_ai/gemini-3.8-flash", + wire="vertex_generate", + override_model="gemini-3.1-pro-preview", + override_map_key="gemini-3.1-pro-preview", + litellm_params=_VERTEX_PARAMS, + ), + _ExtendedSpec( + map_key="gemini-3.1-pro-preview", + litellm_model="vertex_ai/gemini-3.1-pro-preview", + wire="vertex_generate", + override_model="gemini-3.8-flash", + override_map_key="gemini-3.8-flash", + litellm_params=_VERTEX_PARAMS, + ), +) + + def _frontier() -> tuple[FrontierModel, ...]: return tuple( FrontierModel( @@ -174,8 +272,21 @@ def _frontier() -> tuple[FrontierModel, ...]: wire=wire, map_key=map_key, override_model=_OVERRIDE_MODELS[map_key], + override_map_key=_OVERRIDE_MAP_KEYS[_OVERRIDE_MODELS[map_key]], ) for map_key, litellm_model, wire in _FRONTIER_SPECS + ) + tuple( + FrontierModel( + model_name=f"cc-{spec.map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", + litellm_model=spec.litellm_model, + wire=spec.wire, + map_key=spec.map_key, + override_model=spec.override_model, + override_map_key=spec.override_map_key, + base_model=spec.base_model, + litellm_params=spec.litellm_params, + ) + for spec in _EXTENDED_SPECS ) @@ -219,6 +330,24 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), + "azure_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", + } + ), + "bedrock_converse": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "absent_usage", + "tool_call", "image_input", + } + ), + "vertex_generate": frozenset( + { + "cache_read", "reasoning", "audio", "web_search", "response_model", + "absent_usage", "tool_call", "image_input", "prompt_blocked", + } + ), }) CaseName: TypeAlias = Literal[ @@ -270,6 +399,7 @@ class Case: scenario_id=scenario_id, wire=model.wire, usage=self.usage, + model=model.provider_model, output=ScriptedOutput( text=text, response_model=model.override_model if self.response_model_override else None, @@ -296,7 +426,9 @@ _PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tok def _web_search_case(model: FrontierModel) -> Case: - counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + counts_exactly: Final = model.wire in ( + "openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate" + ) return Case( name="web_search", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), @@ -611,12 +743,12 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: wire's normalization: Anthropic folds cache read/write into prompt_tokens, everyone else reports the totals the wire emitted.""" u: Final = case.usage - if model.wire == "anthropic_messages": + if model.wire in ("anthropic_messages", "bedrock_converse"): return ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, u.output_tokens, ) - if model.wire == "gemini_generate": + if model.wire in ("gemini_generate", "vertex_generate"): return ( u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index e1a6c430307..00230fabeba 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -15,10 +15,15 @@ Layout on one port: - ``GET /health`` liveness - ``POST /_scenarios`` register a Scenario JSON, returns its id - ``DELETE /_scenarios/`` remove it +- ``POST /_oauth/token`` fake Google OAuth token endpoint for the + Vertex service-account credential's refresh call - ``POST ///`` provider wire; mount is one of - ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks`` and the - remainder is whatever path the provider client appends (``chat/completions``, - ``responses``, ``v1/messages``, ``models/:generateContent`` ...) + ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks``, ``azure``, + ``bedrock``, ``vertex`` and the remainder is whatever path the provider + client appends (``chat/completions``, ``responses``, ``v1/messages``, + ``models/:generateContent`` ...). Vertex appends ``:generateContent`` / + ``:streamGenerateContent`` to the mount segment itself, and Bedrock Converse + targets ``model//converse`` / ``converse-stream`` A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the @@ -28,15 +33,17 @@ final stream chunk carries usage or the provider reports none. from __future__ import annotations import json +import struct import sys import threading import time +import zlib from collections.abc import Mapping from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from types import MappingProxyType from typing import Final, Literal, TypeAlias -from urllib.parse import urlsplit +from urllib.parse import unquote, urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator @@ -47,6 +54,9 @@ Wire: TypeAlias = Literal[ "gemini_generate", "together_chat", "fireworks_chat", + "azure_chat", + "bedrock_converse", + "vertex_generate", ] WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( @@ -57,6 +67,9 @@ WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( "gemini_generate": "gemini", "together_chat": "together", "fireworks_chat": "fireworks", + "azure_chat": "azure", + "bedrock_converse": "bedrock", + "vertex_generate": "vertex", } ) @@ -69,6 +82,7 @@ _TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( { "openai_responses": frozenset({"incomplete", "unvalidated"}), "gemini_generate": frozenset({"prompt_blocked"}), + "vertex_generate": frozenset({"prompt_blocked"}), } ) @@ -131,6 +145,10 @@ class Scenario(BaseModel): wire: Wire usage: ScriptedUsage output: ScriptedOutput + # The bare provider-facing model name the renderer echoes when the request + # carries no model of its own (Vertex and Bedrock name the model in the URL + # path, not the body). + model: str stream_usage: StreamUsage = "final_chunk" service_tier: ServiceTier | None = None @@ -904,7 +922,208 @@ def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: ) -def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> RenderedResponse: +def _bedrock_usage(u: ScriptedUsage) -> Mapping[str, object]: + # Converse reports uncached input in inputTokens and rides cache reads and + # writes on top-level fields; totalTokens covers every input kind + output. + cache_writes: Final = u.cache_write_5m_tokens + u.cache_write_1h_tokens + return _jobj_opt( + ("inputTokens", u.fresh_input_tokens), + ("outputTokens", u.output_tokens), + ( + "totalTokens", + u.fresh_input_tokens + u.cache_read_tokens + cache_writes + u.output_tokens, + ), + ("cacheReadInputTokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ("cacheWriteInputTokens", cache_writes) if cache_writes else None, + ( + ( + "cacheDetails", + tuple( + _jobj(("inputTokens", count), ("ttl", ttl)) + for count, ttl in ( + (u.cache_write_5m_tokens, "5m"), + (u.cache_write_1h_tokens, "1h"), + ) + if count + ), + ) + if cache_writes + else None + ), + ) + + +def _bedrock_stop_reason(scenario: Scenario) -> str: + if scenario.output.tool_call is not None: + return "tool_use" + return "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + + +def _bedrock_content(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ( + "toolUse", + _jobj( + ("toolUseId", f"tooluse_{scenario.scenario_id}"), + ("name", tool_call.name), + ("input", json.loads(tool_call.arguments)), + ), + ), + ), + ) + return (_jobj(("text", scenario.output.text)),) + + +def _bedrock_body(scenario: Scenario) -> Mapping[str, object]: + return _jobj( + ( + "output", + _jobj( + ( + "message", + _jobj( + ("role", "assistant"), + ("content", _bedrock_content(scenario)), + ), + ), + ), + ), + ("stopReason", _bedrock_stop_reason(scenario)), + ("usage", _bedrock_usage(scenario.usage)), + ("metrics", _jobj(("latencyMs", 42))), + ) + + +def _aws_event_frame(event_type: str, payload: Mapping[str, object]) -> bytes: + """One application/vnd.amazon.eventstream frame: prelude + prelude CRC32 + + headers + JSON payload + message CRC32, matching botocore EventStreamBuffer.""" + try: + from botocore.eventstream import crc32 as _crc32 + except ImportError: + _crc32 = zlib.crc32 + + def _str_header(name: str, value: str) -> bytes: + name_b: Final = name.encode() + value_b: Final = value.encode() + return ( + struct.pack("!B", len(name_b)) + + name_b + + struct.pack("!B", 7) + + struct.pack("!H", len(value_b)) + + value_b + ) + + payload_bytes: Final = json.dumps(payload, default=dict, separators=(",", ":")).encode() + headers_bytes: Final = ( + _str_header(":event-type", event_type) + + _str_header(":content-type", "application/json") + + _str_header(":message-type", "event") + ) + total_length: Final = 12 + len(headers_bytes) + len(payload_bytes) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers_bytes)) + prelude_crc: Final = struct.pack("!I", _crc32(prelude) & 0xFFFFFFFF) + message: Final = prelude + prelude_crc + headers_bytes + payload_bytes + return message + struct.pack("!I", _crc32(message, 0) & 0xFFFFFFFF) + + +def _bedrock_eventstream(scenario: Scenario) -> bytes: + tool_call: Final = scenario.output.tool_call + block_start: Final[tuple[bytes, ...]] = ( + ( + _aws_event_frame( + "contentBlockStart", + _jobj( + ( + "start", + _jobj( + ( + "toolUse", + _jobj( + ("toolUseId", f"tooluse_{scenario.scenario_id}"), + ("name", tool_call.name), + ), + ), + ), + ), + ("contentBlockIndex", 0), + ), + ), + ) + if tool_call is not None + else () + ) + deltas: Final[tuple[bytes, ...]] = ( + tuple( + _aws_event_frame( + "contentBlockDelta", + _jobj( + ("delta", _jobj(("toolUse", _jobj(("input", arguments_slice))))), + ("contentBlockIndex", 0), + ), + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ) + if tool_call is not None + else ( + _aws_event_frame( + "contentBlockDelta", + _jobj( + ("delta", _jobj(("text", scenario.output.text))), + ("contentBlockIndex", 0), + ), + ), + ) + ) + return b"".join( + ( + _aws_event_frame("messageStart", _jobj(("role", "assistant"))), + *block_start, + *deltas, + _aws_event_frame("contentBlockStop", _jobj(("contentBlockIndex", 0))), + _aws_event_frame("messageStop", _jobj(("stopReason", _bedrock_stop_reason(scenario)))), + *( + ( + _aws_event_frame( + "metadata", + _jobj( + ("usage", _bedrock_usage(scenario.usage)), + ("metrics", _jobj(("latencyMs", 42))), + ), + ), + ) + if scenario.stream_usage == "final_chunk" + else () + ), + ) + ) + + +def _render( + scenario: Scenario, *, stream: bool, requested_model: str, path_tail: str +) -> RenderedResponse: + # Azure bridges gpt-5.4+ chat requests carrying function tools onto the + # Responses API, which lands on the same mount at openai/responses. + if scenario.wire == "azure_chat" and path_tail.endswith("openai/responses"): + if stream: + return RenderedResponse( + 200, "text/event-stream", _responses_sse(scenario, requested_model) + ) + return RenderedResponse( + 200, "application/json", _json_bytes(_responses_body(scenario, requested_model)) + ) + if scenario.wire == "bedrock_converse": + if stream: + return RenderedResponse( + 200, "application/vnd.amazon.eventstream", _bedrock_eventstream(scenario) + ) + return RenderedResponse(200, "application/json", _json_bytes(_bedrock_body(scenario))) + if scenario.wire == "vertex_generate": + if stream: + return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model))) if scenario.wire == "anthropic_messages": if stream: return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model)) @@ -917,7 +1136,8 @@ def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> Render if stream: return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model)) return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model))) - # openai_chat, together_chat, fireworks_chat share the OpenAI chat shape. + # openai_chat, together_chat, fireworks_chat and azure_chat share the + # OpenAI chat shape. if stream: return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model)) return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model))) @@ -954,17 +1174,28 @@ def _request_body(body: bytes) -> Mapping[str, object]: return MappingProxyType({}) -def _request_wants_stream(path_tail: str, body: bytes) -> bool: - if ":streamGenerateContent" in path_tail: +def _request_wants_stream(mount_endpoint: str | None, path_tail: str, body: bytes) -> bool: + if mount_endpoint == "streamGenerateContent" or ":streamGenerateContent" in path_tail: + return True + if path_tail.endswith("converse-stream"): return True if not body: return False return _request_body(body).get("stream") is True -def _request_model(body: bytes) -> str: +def _request_model(body: bytes, path_tail: str, scenario: Scenario) -> str: model: Final = _request_body(body).get("model") - return model if isinstance(model, str) else "unknown" + if isinstance(model, str): + return model + # Bedrock Converse names the model in the path: model//converse[-stream]. + if path_tail.startswith("model/"): + path_model: Final = path_tail.split("/", 2)[1] if path_tail.count("/") >= 2 else "" + if path_model: + return unquote(path_model) + # Vertex names it in the URL too, but the mount segment swallowed it when + # the api_base carried a path; fall back to the scenario's declared model. + return scenario.model def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: @@ -972,6 +1203,22 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte segments: Final = tuple(segment for segment in path.split("/") if segment) if method == "GET" and segments == ("health",): return RenderedResponse(200, "application/json", _json_bytes(_jobj(("status", "ok")))) + if segments and segments[0] == "_oauth": + if method == "POST" and segments == ("_oauth", "token"): + return RenderedResponse( + 200, + "application/json", + _json_bytes( + _jobj( + ("access_token", "scripted-token"), + ("token_type", "Bearer"), + ("expires_in", 3600), + ) + ), + ) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", "unknown control route"))) + ) if segments and segments[0] == "_scenarios": if method == "POST" and len(segments) == 1: try: @@ -998,7 +1245,15 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte return RenderedResponse( 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) ) - scenario_id, mount = segments[0], segments[1] + scenario_id: Final = segments[0] + # Vertex builds {api_base}:{endpoint}, so the mount segment can carry a + # :generateContent / :streamGenerateContent suffix. + mount_segment: Final = segments[1] + mount, mount_endpoint = ( + mount_segment.split(":", 1) + if ":" in mount_segment + else (mount_segment, None) + ) found: Final = store.get(scenario_id) if found is None: return RenderedResponse( @@ -1013,7 +1268,12 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte ), ) tail: Final = "/".join(segments[2:]) - return _render(found, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + return _render( + found, + stream=_request_wants_stream(mount_endpoint, tail, body), + requested_model=_request_model(body, tail, found), + path_tail=tail, + ) class _ScriptedHandler(BaseHTTPRequestHandler): diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index 4da7b31a6ef..a36bb1a8662 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -94,6 +94,40 @@ _WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ "fireworks_ai/kimi-k3", ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), ), + "azure_chat": ( + "azure/gpt-5.6", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "bedrock_converse": ( + "anthropic.claude-sonnet-5-v1:0", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + ), + ), + "vertex_generate": ( + "gemini-3.8-flash", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), }) _SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) @@ -141,6 +175,36 @@ _SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( response_model_override=True, ), ), + ( + "vertex_prompt_blocked", + "vertex_generate", + Case( + name="prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "vertex_prompt_blocked_stream", + "vertex_generate", + Case( + name="stream_prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "azure_served_model_override", + "azure_chat", + Case( + name="response_model_override", + usage=_SHAPE_USAGE, + response_model_override=True, + ), + ), ) diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index 68d840870c9..4fba337b701 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -304,6 +304,149 @@ "supports_reasoning": true, "supports_web_search": true }, + "anthropic.claude-sonnet-5-v1:0": { + "cache_creation_input_token_cost": 0.00051, + "cache_creation_input_token_cost_above_1hr": 0.00068, + "cache_read_input_token_cost": 1.7e-05, + "input_cost_per_token": 0.00017, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00034, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00048, + "cache_creation_input_token_cost_above_1hr": 0.00064, + "cache_read_input_token_cost": 1.6e-05, + "input_cost_per_audio_token": 0.00096, + "input_cost_per_token": 0.00016, + "input_cost_per_token_above_200k_tokens": 0.00128, + "input_cost_per_token_flex": 0.00024, + "input_cost_per_token_priority": 0.000272, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00112, + "output_cost_per_reasoning_token": 0.0008, + "output_cost_per_token": 0.00032, + "output_cost_per_token_above_200k_tokens": 0.00144, + "output_cost_per_token_flex": 0.0004, + "output_cost_per_token_priority": 0.000432, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "azure/gpt-5.6": { + "cache_creation_input_token_cost": 0.00044999999999999996, + "cache_creation_input_token_cost_above_1hr": 0.0006000000000000001, + "cache_read_input_token_cost": 1.5e-05, + "input_cost_per_audio_token": 0.0009000000000000001, + "input_cost_per_token": 0.00015000000000000001, + "input_cost_per_token_above_200k_tokens": 0.0012000000000000001, + "input_cost_per_token_flex": 0.000225, + "input_cost_per_token_priority": 0.000255, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0010500000000000002, + "output_cost_per_reasoning_token": 0.00075, + "output_cost_per_token": 0.00030000000000000003, + "output_cost_per_token_above_200k_tokens": 0.00135, + "output_cost_per_token_flex": 0.000375, + "output_cost_per_token_priority": 0.00040499999999999996, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.1e-05, + "input_cost_per_audio_token": 0.00126, + "input_cost_per_token": 0.00021, + "input_cost_per_token_above_200k_tokens": 0.00168, + "input_cost_per_token_flex": 0.000315, + "input_cost_per_token_priority": 0.000357, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00147, + "output_cost_per_reasoning_token": 0.0010500000000000002, + "output_cost_per_token": 0.00042, + "output_cost_per_token_above_200k_tokens": 0.0018900000000000001, + "output_cost_per_token_flex": 0.000525, + "output_cost_per_token_priority": 0.000567, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.8-flash": { + "cache_read_input_token_cost": 2e-05, + "input_cost_per_audio_token": 0.0012, + "input_cost_per_token": 0.0002, + "input_cost_per_token_above_200k_tokens": 0.0016, + "input_cost_per_token_flex": 0.0003, + "input_cost_per_token_priority": 0.00034, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0014000000000000002, + "output_cost_per_reasoning_token": 0.001, + "output_cost_per_token": 0.0004, + "output_cost_per_token_above_200k_tokens": 0.0018000000000000001, + "output_cost_per_token_flex": 0.0005, + "output_cost_per_token_priority": 0.00054, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "meta.llama4-maverick-17b-instruct-v1:0": { + "input_cost_per_token": 0.00019, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00038, + "supports_function_calling": true + }, "together_ai/moonshotai/Kimi-K3": { "cache_creation_input_token_cost": 0.00030000000000000003, "cache_creation_input_token_cost_above_1hr": 0.0004, @@ -363,5 +506,20 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_web_search": true + }, + "us.anthropic.claude-opus-5-v1:0": { + "cache_creation_input_token_cost": 0.0005400000000000001, + "cache_creation_input_token_cost_above_1hr": 0.00072, + "cache_read_input_token_cost": 1.8e-05, + "input_cost_per_token": 0.00018, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00036000000000000004, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true } } diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..d96478de1c4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1000,6 +1000,7 @@ class ModelInfoBody(BaseModel): access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None + base_model: str | None = None class ModelNewBody(BaseModel): From 2466975d290576de9e89a1d5d69c9ca9a6aab1ab Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:57:23 +0000 Subject: [PATCH 040/349] test(e2e): clean cost map decimals and simplify scripted wire helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 52 ++-- .../e2e/cost_calculation/scripted_provider.py | 39 ++- tests/e2e/cost_map.json | 270 +++++++++--------- 3 files changed, 177 insertions(+), 184 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 8c6db7c0010..3f3e9fd9243 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -11,6 +11,7 @@ Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). from __future__ import annotations +import functools import importlib.util import json import sys @@ -21,6 +22,8 @@ from types import ModuleType from typing import Final, Protocol, cast import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa from cost_matrix import Case, FrontierModel from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE @@ -112,33 +115,25 @@ def client() -> CostCalcClient: return CostCalcClient(proxy=proxy) -_vertex_key_pem: str | None = None +@functools.cache +def _vertex_private_key_pem() -> str: + return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() def _vertex_service_account_json() -> str: """A service-account credential JSON whose token_uri is the sidecar's /_oauth/token route: the proxy's google-auth refresh then gets a scripted - access token without touching Google. One generated RSA key per process.""" - global _vertex_key_pem # mutable-ok: session-scoped key generation cached for reuse - if _vertex_key_pem is None: - from cryptography.hazmat.primitives import serialization - from cryptography.hazmat.primitives.asymmetric import rsa - - _vertex_key_pem = ( - rsa.generate_private_key(public_exponent=65537, key_size=2048) - .private_bytes( - serialization.Encoding.PEM, - serialization.PrivateFormat.PKCS8, - serialization.NoEncryption(), - ) - .decode() - ) + access token without touching Google.""" return json.dumps( { "type": "service_account", "project_id": "cc-scripted-project", "private_key_id": "scripted", - "private_key": _vertex_key_pem, + "private_key": _vertex_private_key_pem(), "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com", "client_id": "0", "auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize", @@ -162,20 +157,21 @@ def register_scenario_deployment( handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) model_name: Final = f"{model.model_name}-{marker}" - extra_params: Final[dict[str, str]] = dict(model.litellm_params) - if model.wire == "vertex_generate": - extra_params["vertex_credentials"] = _vertex_service_account_json() + params: Final = { + "model": model.litellm_model, + "api_key": model.api_key, + "api_base": handle.api_base(), + **model.litellm_params, + **( + {"vertex_credentials": _vertex_service_account_json()} + if model.wire == "vertex_generate" + else {} + ), + } model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, - litellm_params=LiteLLMParamsBody.model_validate( - { - "model": model.litellm_model, - "api_key": model.api_key, - "api_base": handle.api_base(), - **extra_params, - } - ), + litellm_params=LiteLLMParamsBody.model_validate(params), model_info=ModelInfoBody(base_model=model.base_model), ) ) diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 00230fabeba..982132ed8df 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -997,36 +997,33 @@ def _bedrock_body(scenario: Scenario) -> Mapping[str, object]: ) +def _aws_str_header(name: str, value: str) -> bytes: + """One eventstream header: 1-byte name len + name + type-7 marker + value.""" + name_b: Final = name.encode() + value_b: Final = value.encode() + return ( + struct.pack("!B", len(name_b)) + + name_b + + struct.pack("!B", 7) + + struct.pack("!H", len(value_b)) + + value_b + ) + + def _aws_event_frame(event_type: str, payload: Mapping[str, object]) -> bytes: """One application/vnd.amazon.eventstream frame: prelude + prelude CRC32 + headers + JSON payload + message CRC32, matching botocore EventStreamBuffer.""" - try: - from botocore.eventstream import crc32 as _crc32 - except ImportError: - _crc32 = zlib.crc32 - - def _str_header(name: str, value: str) -> bytes: - name_b: Final = name.encode() - value_b: Final = value.encode() - return ( - struct.pack("!B", len(name_b)) - + name_b - + struct.pack("!B", 7) - + struct.pack("!H", len(value_b)) - + value_b - ) - payload_bytes: Final = json.dumps(payload, default=dict, separators=(",", ":")).encode() headers_bytes: Final = ( - _str_header(":event-type", event_type) - + _str_header(":content-type", "application/json") - + _str_header(":message-type", "event") + _aws_str_header(":event-type", event_type) + + _aws_str_header(":content-type", "application/json") + + _aws_str_header(":message-type", "event") ) total_length: Final = 12 + len(headers_bytes) + len(payload_bytes) + 4 prelude: Final = struct.pack("!II", total_length, len(headers_bytes)) - prelude_crc: Final = struct.pack("!I", _crc32(prelude) & 0xFFFFFFFF) + prelude_crc: Final = struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) message: Final = prelude + prelude_crc + headers_bytes + payload_bytes - return message + struct.pack("!I", _crc32(message, 0) & 0xFFFFFFFF) + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) def _bedrock_eventstream(scenario: Scenario) -> bytes: diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index 4fba337b701..85cd5ade3d5 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -1,4 +1,79 @@ { + "anthropic.claude-sonnet-5-v1:0": { + "cache_creation_input_token_cost": 0.00051, + "cache_creation_input_token_cost_above_1hr": 0.00068, + "cache_read_input_token_cost": 1.7e-05, + "input_cost_per_token": 0.00017, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00034, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00048, + "cache_creation_input_token_cost_above_1hr": 0.00064, + "cache_read_input_token_cost": 1.6e-05, + "input_cost_per_audio_token": 0.00096, + "input_cost_per_token": 0.00016, + "input_cost_per_token_above_200k_tokens": 0.00128, + "input_cost_per_token_flex": 0.00024, + "input_cost_per_token_priority": 0.000272, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00112, + "output_cost_per_reasoning_token": 0.0008, + "output_cost_per_token": 0.00032, + "output_cost_per_token_above_200k_tokens": 0.00144, + "output_cost_per_token_flex": 0.0004, + "output_cost_per_token_priority": 0.000432, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "azure/gpt-5.6": { + "cache_creation_input_token_cost": 0.00045, + "cache_creation_input_token_cost_above_1hr": 0.0006, + "cache_read_input_token_cost": 1.5e-05, + "input_cost_per_audio_token": 0.0009, + "input_cost_per_token": 0.00015, + "input_cost_per_token_above_200k_tokens": 0.0012, + "input_cost_per_token_flex": 0.000225, + "input_cost_per_token_priority": 0.000255, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00105, + "output_cost_per_reasoning_token": 0.00075, + "output_cost_per_token": 0.0003, + "output_cost_per_token_above_200k_tokens": 0.00135, + "output_cost_per_token_flex": 0.000375, + "output_cost_per_token_priority": 0.000405, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, "claude-haiku-4-5": { "cache_creation_input_token_cost": 0.00021, "cache_creation_input_token_cost_above_1hr": 0.00028000000000000003, @@ -134,6 +209,64 @@ "supports_reasoning": true, "supports_web_search": true }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.1e-05, + "input_cost_per_audio_token": 0.00126, + "input_cost_per_token": 0.00021, + "input_cost_per_token_above_200k_tokens": 0.00168, + "input_cost_per_token_flex": 0.000315, + "input_cost_per_token_priority": 0.000357, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00147, + "output_cost_per_reasoning_token": 0.00105, + "output_cost_per_token": 0.00042, + "output_cost_per_token_above_200k_tokens": 0.00189, + "output_cost_per_token_flex": 0.000525, + "output_cost_per_token_priority": 0.000567, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.8-flash": { + "cache_read_input_token_cost": 2e-05, + "input_cost_per_audio_token": 0.0012, + "input_cost_per_token": 0.0002, + "input_cost_per_token_above_200k_tokens": 0.0016, + "input_cost_per_token_flex": 0.0003, + "input_cost_per_token_priority": 0.00034, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0014, + "output_cost_per_reasoning_token": 0.001, + "output_cost_per_token": 0.0004, + "output_cost_per_token_above_200k_tokens": 0.0018, + "output_cost_per_token_flex": 0.0005, + "output_cost_per_token_priority": 0.00054, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 9e-06, "input_cost_per_audio_token": 0.00054, @@ -304,139 +437,6 @@ "supports_reasoning": true, "supports_web_search": true }, - "anthropic.claude-sonnet-5-v1:0": { - "cache_creation_input_token_cost": 0.00051, - "cache_creation_input_token_cost_above_1hr": 0.00068, - "cache_read_input_token_cost": 1.7e-05, - "input_cost_per_token": 0.00017, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 0.00034, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true - }, - "azure/gpt-5.4-mini": { - "cache_creation_input_token_cost": 0.00048, - "cache_creation_input_token_cost_above_1hr": 0.00064, - "cache_read_input_token_cost": 1.6e-05, - "input_cost_per_audio_token": 0.00096, - "input_cost_per_token": 0.00016, - "input_cost_per_token_above_200k_tokens": 0.00128, - "input_cost_per_token_flex": 0.00024, - "input_cost_per_token_priority": 0.000272, - "litellm_provider": "azure", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.00112, - "output_cost_per_reasoning_token": 0.0008, - "output_cost_per_token": 0.00032, - "output_cost_per_token_above_200k_tokens": 0.00144, - "output_cost_per_token_flex": 0.0004, - "output_cost_per_token_priority": 0.000432, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true - }, - "azure/gpt-5.6": { - "cache_creation_input_token_cost": 0.00044999999999999996, - "cache_creation_input_token_cost_above_1hr": 0.0006000000000000001, - "cache_read_input_token_cost": 1.5e-05, - "input_cost_per_audio_token": 0.0009000000000000001, - "input_cost_per_token": 0.00015000000000000001, - "input_cost_per_token_above_200k_tokens": 0.0012000000000000001, - "input_cost_per_token_flex": 0.000225, - "input_cost_per_token_priority": 0.000255, - "litellm_provider": "azure", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.0010500000000000002, - "output_cost_per_reasoning_token": 0.00075, - "output_cost_per_token": 0.00030000000000000003, - "output_cost_per_token_above_200k_tokens": 0.00135, - "output_cost_per_token_flex": 0.000375, - "output_cost_per_token_priority": 0.00040499999999999996, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true - }, - "gemini-3.1-pro-preview": { - "cache_read_input_token_cost": 2.1e-05, - "input_cost_per_audio_token": 0.00126, - "input_cost_per_token": 0.00021, - "input_cost_per_token_above_200k_tokens": 0.00168, - "input_cost_per_token_flex": 0.000315, - "input_cost_per_token_priority": 0.000357, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.00147, - "output_cost_per_reasoning_token": 0.0010500000000000002, - "output_cost_per_token": 0.00042, - "output_cost_per_token_above_200k_tokens": 0.0018900000000000001, - "output_cost_per_token_flex": 0.000525, - "output_cost_per_token_priority": 0.000567, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true, - "web_search_billing_unit": "per_query" - }, - "gemini-3.8-flash": { - "cache_read_input_token_cost": 2e-05, - "input_cost_per_audio_token": 0.0012, - "input_cost_per_token": 0.0002, - "input_cost_per_token_above_200k_tokens": 0.0016, - "input_cost_per_token_flex": 0.0003, - "input_cost_per_token_priority": 0.00034, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.0014000000000000002, - "output_cost_per_reasoning_token": 0.001, - "output_cost_per_token": 0.0004, - "output_cost_per_token_above_200k_tokens": 0.0018000000000000001, - "output_cost_per_token_flex": 0.0005, - "output_cost_per_token_priority": 0.00054, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true, - "web_search_billing_unit": "per_query" - }, "meta.llama4-maverick-17b-instruct-v1:0": { "input_cost_per_token": 0.00019, "litellm_provider": "bedrock_converse", @@ -508,7 +508,7 @@ "supports_web_search": true }, "us.anthropic.claude-opus-5-v1:0": { - "cache_creation_input_token_cost": 0.0005400000000000001, + "cache_creation_input_token_cost": 0.00054, "cache_creation_input_token_cost_above_1hr": 0.00072, "cache_read_input_token_cost": 1.8e-05, "input_cost_per_token": 0.00018, @@ -517,7 +517,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 0.00036000000000000004, + "output_cost_per_token": 0.00036, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true From 813d96f26ea6780bacb4b5ad562f1aecc3cb5069 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:18:26 +0000 Subject: [PATCH 041/349] fix(e2e): resolve remaining merge markers in e2e_config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/e2e_config.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index b34eadd8744..e19cfaa684f 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,7 +145,6 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" -<<<<<<< HEAD # The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL # pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a # scripted-provider sidecar; deselected unless the opt-in env var is set. @@ -162,9 +161,6 @@ SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get( SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get( "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL ).rstrip("/") -||||||| 930ec9643a -======= -CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) From bdfff602fb0325f88768f7c4411cce921ab28fcb Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:47:55 +0000 Subject: [PATCH 042/349] test(e2e): drive the cost matrix from cases.json and expected.json goldens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/cost_calculation/cases.json | 231 ++ tests/e2e/cost_calculation/conftest.py | 10 +- tests/e2e/cost_calculation/cost_matrix.py | 788 ++----- tests/e2e/cost_calculation/expected.json | 2004 +++++++++++++++++ .../e2e/cost_calculation/generate_expected.py | 189 ++ .../e2e/cost_calculation/test_matrix_data.py | 64 + .../test_token_pricing_e2e.py | 62 +- .../cost_calculation/test_wire_formats_e2e.py | 368 --- 9 files changed, 2761 insertions(+), 957 deletions(-) create mode 100644 tests/e2e/cost_calculation/cases.json create mode 100644 tests/e2e/cost_calculation/expected.json create mode 100644 tests/e2e/cost_calculation/generate_expected.py create mode 100644 tests/e2e/cost_calculation/test_matrix_data.py delete mode 100644 tests/e2e/cost_calculation/test_wire_formats_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 49cfc29aa17..707d35b4aa6 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json new file mode 100644 index 00000000000..e898557ea35 --- /dev/null +++ b/tests/e2e/cost_calculation/cases.json @@ -0,0 +1,231 @@ +{ + "deployments": [ + { + "map_key": "azure/gpt-5.4-mini", + "litellm_model": "azure/cc-pinned-deployment", + "base_model": "azure/gpt-5.4-mini" + } + ], + "cases": [ + { + "name": "basic", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40} + }, + { + "name": "cache_read", + "usage": {"fresh_input_tokens": 100, "cache_read_tokens": 50, "output_tokens": 30}, + "requires_rates": ["cache_read_input_token_cost"], + "requires_caps": ["cache_read"] + }, + { + "name": "cache_write_5m", + "usage": {"fresh_input_tokens": 90, "cache_write_5m_tokens": 60, "output_tokens": 30}, + "requires_rates": ["cache_creation_input_token_cost"], + "requires_caps": ["cache_write_5m"] + }, + { + "name": "cache_write_1h", + "usage": {"fresh_input_tokens": 90, "cache_write_5m_tokens": 20, "cache_write_1h_tokens": 40, "output_tokens": 30}, + "requires_rates": ["cache_creation_input_token_cost_above_1hr", "cache_creation_input_token_cost"], + "requires_caps": ["cache_write_1h"] + }, + { + "name": "reasoning", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "reasoning_tokens": 70}, + "requires_rates": ["output_cost_per_reasoning_token"], + "requires_caps": ["reasoning"] + }, + { + "name": "audio", + "usage": {"fresh_input_tokens": 100, "audio_input_tokens": 25, "output_tokens": 30, "audio_output_tokens": 15}, + "requires_rates": ["input_cost_per_audio_token", "output_cost_per_audio_token"], + "requires_caps": ["audio"] + }, + { + "name": "tiered", + "usage": {"fresh_input_tokens": 200001, "output_tokens": 30}, + "requires_rates": ["input_cost_per_token_above_200k_tokens", "output_cost_per_token_above_200k_tokens"] + }, + { + "name": "service_tier_flex", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "service_tier": "flex", + "requires_rates": ["input_cost_per_token_flex", "output_cost_per_token_flex"] + }, + { + "name": "service_tier_priority", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "service_tier": "priority", + "requires_rates": ["input_cost_per_token_priority", "output_cost_per_token_priority"] + }, + { + "name": "web_search", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3}, + "requires_rates": ["search_context_cost_per_query"], + "requires_caps": ["web_search"] + }, + { + "name": "stream", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true + }, + { + "name": "stream_no_usage", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "exact_spend": false, + "requires_caps": ["absent_usage"] + }, + { + "name": "response_model_override", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "response_model_override": true, + "requires_caps": ["response_model"] + }, + { + "name": "stream_response_model_override", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "response_model_override": true, + "requires_caps": ["response_model"] + }, + { + "name": "tool_call", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "tool_call": true, + "requires_caps": ["tool_call"] + }, + { + "name": "stream_tool_call", + "usage": {"fresh_input_tokens": 80, "output_tokens": 25}, + "stream": true, + "tool_call": true, + "requires_caps": ["tool_call"] + }, + { + "name": "stream_no_usage_tool_call", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "tool_call": true, + "exact_spend": false, + "requires_caps": ["absent_usage", "tool_call"] + }, + { + "name": "stream_no_usage_image_input", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "image_input": true, + "exact_spend": false, + "requires_caps": ["absent_usage", "image_input"] + }, + { + "name": "stream_incomplete", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "terminal": "incomplete", + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_no_usage_incomplete", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "terminal": "incomplete", + "exact_spend": false, + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_unvalidated", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "terminal": "unvalidated", + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_no_usage_unvalidated", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "terminal": "unvalidated", + "exact_spend": false, + "requires_caps": ["responses_terminal"] + }, + { + "name": "prompt_blocked", + "usage": {"fresh_input_tokens": 1000, "output_tokens": 0}, + "terminal": "prompt_blocked", + "response_model_override": true, + "requires_caps": ["prompt_blocked"] + }, + { + "name": "stream_prompt_blocked", + "usage": {"fresh_input_tokens": 1000, "output_tokens": 0}, + "stream": true, + "terminal": "prompt_blocked", + "response_model_override": true, + "requires_caps": ["prompt_blocked"] + }, + { + "name": "all_components_chat", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25, + "reasoning_tokens": 15, + "audio_input_tokens": 5, + "audio_output_tokens": 3 + }, + "wires": ["openai_chat", "azure_chat", "together_chat"] + }, + { + "name": "all_components_fireworks", + "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, + "wires": ["fireworks_chat"] + }, + { + "name": "all_components_anthropic", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25 + }, + "wires": ["anthropic_messages", "bedrock_converse"] + }, + { + "name": "all_components_anthropic_stream", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25 + }, + "stream": true, + "wires": ["anthropic_messages"] + }, + { + "name": "all_components_gemini", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "output_tokens": 25, + "reasoning_tokens": 15, + "audio_input_tokens": 5, + "audio_output_tokens": 3 + }, + "wires": ["gemini_generate", "vertex_generate"] + }, + { + "name": "all_components_responses", + "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, + "wires": ["openai_responses"] + } + ] +} diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 3f3e9fd9243..3de9786854e 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -1,10 +1,12 @@ """Cost-calculation suite fixtures. Runs against a dedicated proxy whose whole model cost map is the test-owned -``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL), so every deployment -bills at rates the test asserts literal arithmetic on. Provider calls are -answered by the scripted-provider sidecar (``scripted_provider.py``), registered -per scenario over its control API. +``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL); every map entry is a +deployment under test, the request shapes live in ``cases.json``, and the +asserted goldens live in ``expected.json`` (regenerate proposals with +``generate_expected.py``). Provider calls are answered by the +scripted-provider sidecar (``scripted_provider.py``), registered per scenario +over its control API. Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). """ diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 37495f37b0b..b03c851d208 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -1,18 +1,16 @@ -"""The cost-calculation matrix: frontier model set, the pricing-component cases -each model runs, and the expected-cost arithmetic. +"""The cost-calculation matrix: the model set derived from the test cost map, +the request/response cases from ``cases.json``, and the loaders both use. -Rates come from ``tests/e2e/cost_map.json``, which the proxy under test loads as -its ENTIRE model cost map (LITELLM_MODEL_COST_MAP_URL), so an entry's rates are -exactly what the proxy bills and nothing in the suite depends on the bundled -map. Each model's rates are a distinct multiple of a shared base set, so a -component billed at the wrong model's rate (or the wrong case's rate) can never -coincidentally match. - -Case applicability is pricing-field-gated AND wire-gated: a case runs for a -model only when the entry carries the rate the case exercises and the wire can -report the token kind that rate prices. When the wire cannot report a kind -(e.g. Anthropic has no reasoning-token field, Responses reports no cache -creation), the case is absent from the matrix rather than silently zero. +Three data files drive the suite; nothing in Python lists models or cases: +- ``tests/e2e/cost_map.json`` is the proxy's ENTIRE model cost map + (LITELLM_MODEL_COST_MAP_URL); every entry becomes a deployment under test. +- ``tests/e2e/cost_calculation/cases.json`` is the case list; each case runs + for a model when the entry carries the rates it exercises (``requires_rates``) + and the wire can report the token kinds involved (``requires_caps`` / + ``wires``). +- ``tests/e2e/cost_calculation/expected.json`` holds the reviewed goldens; the + tests assert them verbatim and never compute a price themselves. The rate + arithmetic that proposes goldens lives in ``generate_expected.py``, not here. """ from __future__ import annotations @@ -26,13 +24,15 @@ from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from types import MappingProxyType -from typing import Final, Literal, TypeAlias +from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" +CASES_PATH: Final = Path(__file__).resolve().parent / "cases.json" +EXPECTED_PATH: Final = Path(__file__).resolve().parent / "expected.json" class SearchContextCostPerQuery(BaseModel): @@ -77,119 +77,90 @@ _COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( TIER_THRESHOLD_TOKENS: Final = 200_000 -@dataclass(frozen=True, slots=True) -class FrontierModel: - """One deployment under test: the model_name the suite registers, the - provider-prefixed litellm model string, the wire the scripted upstream - speaks, its cost-map key, and the sibling map model the response_model - override case reports.""" +class DeploymentSpec(BaseModel): + """A deployment-level fact from cases.json: when a map key needs a + registered deployment name that is not its provider model (or a + model_info.base_model pin), the matrix uses these instead of the defaults.""" + + model_config = ConfigDict(frozen=True) - model_name: str - litellm_model: str - wire: Wire map_key: str - override_model: str | None = None - override_map_key: str | None = None - # Registered as model_info.base_model; when set, the provider-reported - # model loses to it and every case bills at this deployment's own rates. + litellm_model: str | None = None base_model: str | None = None - # Extra litellm_params merged into the /model/new registration (api_version, - # aws_* credentials, vertex_* auth). - litellm_params: Mapping[str, str] = MappingProxyType({}) - - @property - def rates(self) -> CostMapEntry: - return _COST_MAP[self.map_key] - - @property - def override_rates(self) -> CostMapEntry: - if self.base_model is not None or self.override_map_key is None: - return self.rates - return _COST_MAP[self.override_map_key] - - @property - def provider_model(self) -> str: - """The bare provider-facing model name: litellm_model minus the provider - prefix and any routing segment (converse/, responses/).""" - tail: Final = self.litellm_model.split("/")[1:] - return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) - - @property - def provider(self) -> str: - return self.rates.litellm_provider - - @property - def api_key(self) -> str: - # The scripted upstream ignores auth; a fixed bogus key proves the suite - # spends zero real provider calls. - return "sk-scripted-provider" -# Response-model override targets: emit a sibling's bare provider-facing name so -# the biller's provider-prefixed lookup lands on that sibling's map key. -_OVERRIDE_MODELS: Final[Mapping[str, str]] = MappingProxyType({ - "gpt-5.6": "gpt-5.4-mini", - "gpt-5.5-pro": "gpt-5.3-codex", - "gpt-5.3-codex": "gpt-5.5-pro", - "gpt-5.4-mini": "gpt-5.6", - "claude-opus-5": "claude-sonnet-5", - "claude-sonnet-5": "claude-opus-5", - "claude-haiku-4-5": "claude-sonnet-5", - "gemini/gemini-3.8-flash": "gemini-3.1-pro-preview", - "gemini/gemini-3.1-pro-preview": "gemini-3.8-flash", - "together_ai/moonshotai/Kimi-K3": "zai-org/GLM-5.3", - "together_ai/zai-org/GLM-5.3": "moonshotai/Kimi-K3", - "fireworks_ai/kimi-k3": "qwen3p8-max", - "fireworks_ai/qwen3p8-max": "kimi-k3", - "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", -}) +class Case(BaseModel): + """One request/response shape from cases.json; gated onto a model by + ``requires_rates`` (entry must carry each rate field), ``requires_caps`` + (the wire must report the token kind) and ``wires`` (shape is wire-specific).""" -_OVERRIDE_MAP_KEYS: Final[Mapping[str, str]] = MappingProxyType({ - "gpt-5.4-mini": "gpt-5.4-mini", - "gpt-5.6": "gpt-5.6", - "gpt-5.3-codex": "gpt-5.3-codex", - "gpt-5.5-pro": "gpt-5.5-pro", - "claude-sonnet-5": "claude-sonnet-5", - "claude-opus-5": "claude-opus-5", - "gemini-3.1-pro-preview": "gemini/gemini-3.1-pro-preview", - "gemini-3.8-flash": "gemini/gemini-3.8-flash", - "zai-org/GLM-5.3": "together_ai/zai-org/GLM-5.3", - "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", - "qwen3p8-max": "fireworks_ai/qwen3p8-max", - "kimi-k3": "fireworks_ai/kimi-k3", -}) + model_config = ConfigDict(frozen=True) + + name: str + usage: ScriptedUsage + stream: bool = False + stream_usage: Literal["final_chunk", "absent"] = "final_chunk" + service_tier: Literal["flex", "priority"] | None = None + response_model_override: bool = False + exact_spend: bool = True + tool_call: bool = False + image_input: bool = False + terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" + requires_rates: tuple[str, ...] = () + requires_caps: tuple[str, ...] = () + wires: tuple[Wire, ...] | None = None + + def applies_to(self, model: FrontierModel) -> bool: + if self.wires is not None and model.wire not in self.wires: + return False + caps: Final = _WIRE_CAPS[model.wire] + if not frozenset(self.requires_caps) <= caps: + return False + return all( + getattr(model.rates, field, None) is not None for field in self.requires_rates + ) + + def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: + return Scenario( + scenario_id=scenario_id, + wire=model.wire, + usage=self.usage, + model=model.provider_model, + output=ScriptedOutput( + text=text, + response_model=model.override_model if self.response_model_override else None, + tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) + if self.tool_call + else None, + terminal=self.terminal, + ), + stream_usage=self.stream_usage, + service_tier=self.service_tier, + ) -_FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( - ("gpt-5.6", "openai/gpt-5.6", "openai_chat"), - ("gpt-5.5-pro", "openai/gpt-5.5-pro", "openai_responses"), - ("gpt-5.3-codex", "openai/gpt-5.3-codex", "openai_responses"), - ("gpt-5.4-mini", "openai/gpt-5.4-mini", "openai_chat"), - ("claude-opus-5", "anthropic/claude-opus-5", "anthropic_messages"), - ("claude-sonnet-5", "anthropic/claude-sonnet-5", "anthropic_messages"), - ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "anthropic_messages"), - ("gemini/gemini-3.8-flash", "gemini/gemini-3.8-flash", "gemini_generate"), - ("gemini/gemini-3.1-pro-preview", "gemini/gemini-3.1-pro-preview", "gemini_generate"), - ("together_ai/moonshotai/Kimi-K3", "together_ai/moonshotai/Kimi-K3", "together_chat"), - ("together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3", "together_chat"), - ("fireworks_ai/kimi-k3", "fireworks_ai/kimi-k3", "fireworks_chat"), - ("fireworks_ai/qwen3p8-max", "fireworks_ai/qwen3p8-max", "fireworks_chat"), - ("fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/deepseek-v4p1-flash", "fireworks_chat"), +class _CasesFile(BaseModel): + model_config = ConfigDict(frozen=True) + + deployments: tuple[DeploymentSpec, ...] = () + cases: tuple[Case, ...] = () + + +_CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) +CASES: Final[tuple[Case, ...]] = _CASES_FILE.cases +_DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( + {spec.map_key: spec for spec in _CASES_FILE.deployments} ) @dataclass(frozen=True, slots=True) -class _ExtendedSpec: - """A frontier entry whose override target, model_info.base_model or extra - litellm_params can't be derived from the map key alone.""" +class _ProviderWiring: + """How a (litellm_provider, mode) pair maps to a sidecar wire, the provider + prefix on the registered litellm model string, and extra litellm_params.""" - map_key: str - litellm_model: str wire: Wire - override_model: str | None = None - override_map_key: str | None = None - base_model: str | None = None - litellm_params: Mapping[str, str] = MappingProxyType({}) + model_prefix: str | None + litellm_params: Mapping[str, str] _AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"}) @@ -207,87 +178,134 @@ _VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType( } ) -_EXTENDED_SPECS: Final[tuple[_ExtendedSpec, ...]] = ( - _ExtendedSpec( - map_key="azure/gpt-5.6", - litellm_model="azure/gpt-5.6", - wire="azure_chat", - override_model="gpt-5.4-mini", - override_map_key="azure/gpt-5.4-mini", - litellm_params=_AZURE_PARAMS, - ), - _ExtendedSpec( - # Deployment name is not a model; base_model pins billing so the - # response's model field loses, proving base_model wins. - map_key="azure/gpt-5.4-mini", - litellm_model="azure/cc-pinned-deployment", - wire="azure_chat", - override_model="gpt-5.6", - override_map_key="azure/gpt-5.6", - base_model="azure/gpt-5.4-mini", - litellm_params=_AZURE_PARAMS, - ), - _ExtendedSpec( - map_key="anthropic.claude-sonnet-5-v1:0", - litellm_model="bedrock/converse/anthropic.claude-sonnet-5-v1:0", - wire="bedrock_converse", - litellm_params=_BEDROCK_PARAMS, - ), - _ExtendedSpec( - map_key="us.anthropic.claude-opus-5-v1:0", - litellm_model="bedrock/converse/us.anthropic.claude-opus-5-v1:0", - wire="bedrock_converse", - litellm_params=_BEDROCK_PARAMS, - ), - _ExtendedSpec( - map_key="meta.llama4-maverick-17b-instruct-v1:0", - litellm_model="bedrock/converse/meta.llama4-maverick-17b-instruct-v1:0", - wire="bedrock_converse", - litellm_params=_BEDROCK_PARAMS, - ), - _ExtendedSpec( - map_key="gemini-3.8-flash", - litellm_model="vertex_ai/gemini-3.8-flash", - wire="vertex_generate", - override_model="gemini-3.1-pro-preview", - override_map_key="gemini-3.1-pro-preview", - litellm_params=_VERTEX_PARAMS, - ), - _ExtendedSpec( - map_key="gemini-3.1-pro-preview", - litellm_model="vertex_ai/gemini-3.1-pro-preview", - wire="vertex_generate", - override_model="gemini-3.8-flash", - override_map_key="gemini-3.8-flash", - litellm_params=_VERTEX_PARAMS, - ), +_PROVIDER_WIRING: Final[Mapping[tuple[str, str], _ProviderWiring]] = MappingProxyType( + { + ("openai", "chat"): _ProviderWiring("openai_chat", "openai", MappingProxyType({})), + ("openai", "responses"): _ProviderWiring( + "openai_responses", "openai", MappingProxyType({}) + ), + ("anthropic", "chat"): _ProviderWiring( + "anthropic_messages", "anthropic", MappingProxyType({}) + ), + ("gemini", "chat"): _ProviderWiring("gemini_generate", None, MappingProxyType({})), + ("together_ai", "chat"): _ProviderWiring("together_chat", None, MappingProxyType({})), + ("fireworks_ai", "chat"): _ProviderWiring("fireworks_chat", None, MappingProxyType({})), + ("azure", "chat"): _ProviderWiring("azure_chat", None, _AZURE_PARAMS), + ("bedrock_converse", "chat"): _ProviderWiring( + "bedrock_converse", "bedrock/converse", _BEDROCK_PARAMS + ), + ("vertex_ai-language-models", "chat"): _ProviderWiring( + "vertex_generate", "vertex_ai", _VERTEX_PARAMS + ), + } ) +@dataclass(frozen=True, slots=True) +class FrontierModel: + """One deployment under test, derived from a cost-map entry: the model_name + the suite registers, the provider-prefixed litellm model string, the wire + the scripted upstream speaks, and the sibling map model the response_model + override case reports.""" + + model_name: str + litellm_model: str + wire: Wire + map_key: str + override_model: str | None = None + override_map_key: str | None = None + # Registered as model_info.base_model; when set, the provider-reported + # model loses to it and every case bills at this deployment's own rates. + base_model: str | None = None + litellm_params: Mapping[str, str] = MappingProxyType({}) + + @property + def rates(self) -> CostMapEntry: + return _COST_MAP[self.map_key] + + @property + def override_rates(self) -> CostMapEntry: + if self.base_model is not None or self.override_map_key is None: + return self.rates + return _COST_MAP[self.override_map_key] + + @property + def provider_model(self) -> str: + """The bare provider-facing model name: litellm_model minus the provider + prefix and any routing segment (converse/, responses/).""" + return _provider_model(self.litellm_model) + + @property + def provider(self) -> str: + return self.rates.litellm_provider + + @property + def api_key(self) -> str: + # The scripted upstream ignores auth; a fixed bogus key proves the suite + # spends zero real provider calls. + return "sk-scripted-provider" + + +def _provider_model(litellm_model: str) -> str: + tail: Final = litellm_model.split("/")[1:] + return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) + + +def _litellm_model_for(map_key: str, wiring: _ProviderWiring) -> str: + if wiring.model_prefix is None: + return map_key + if map_key.startswith(f"{wiring.model_prefix}/"): + return map_key + return f"{wiring.model_prefix}/{map_key}" + + def _frontier() -> tuple[FrontierModel, ...]: - return tuple( - FrontierModel( - model_name=f"cc-{map_key.replace('/', '-').lower()}", - litellm_model=litellm_model, - wire=wire, - map_key=map_key, - override_model=_OVERRIDE_MODELS[map_key], - override_map_key=_OVERRIDE_MAP_KEYS[_OVERRIDE_MODELS[map_key]], - ) - for map_key, litellm_model, wire in _FRONTIER_SPECS - ) + tuple( - FrontierModel( - model_name=f"cc-{spec.map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", - litellm_model=spec.litellm_model, - wire=spec.wire, - map_key=spec.map_key, - override_model=spec.override_model, - override_map_key=spec.override_map_key, - base_model=spec.base_model, - litellm_params=spec.litellm_params, - ) - for spec in _EXTENDED_SPECS + groups: Final[Mapping[tuple[str, str], tuple[str, ...]]] = MappingProxyType( + { + pair: tuple(sorted(k for k, e in _COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) + for pair in {(e.litellm_provider, e.mode) for e in _COST_MAP.values()} + } ) + models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple + for map_key in sorted(_COST_MAP): + entry: Final = _COST_MAP[map_key] + pair: Final = (entry.litellm_provider, entry.mode) + wiring: Final = _PROVIDER_WIRING.get(pair) + if wiring is None: + raise ValueError( + f"cost_map entry {map_key} has no wiring for " + f"(litellm_provider={pair[0]}, mode={pair[1]}); add a " + f"_ProviderWiring row in cost_matrix.py" + ) + siblings: Final = groups[pair] + override_key: Final = ( + siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None + ) + override_litellm: Final = ( + _litellm_model_for(override_key, wiring) if override_key is not None else None + ) + deployment: Final = _DEPLOYMENTS.get(map_key) + models.append( + FrontierModel( + model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", + litellm_model=( + deployment.litellm_model + if deployment is not None and deployment.litellm_model is not None + else _litellm_model_for(map_key, wiring) + ), + wire=wiring.wire, + map_key=map_key, + override_model=( + _provider_model(override_litellm) + if override_litellm is not None + else None + ), + override_map_key=override_key, + base_model=deployment.base_model if deployment is not None else None, + litellm_params=wiring.litellm_params, + ) + ) + return tuple(models) FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() @@ -350,71 +368,6 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ ), }) -CaseName: TypeAlias = Literal[ - "basic", - "cache_read", - "cache_write_5m", - "cache_write_1h", - "reasoning", - "audio", - "tiered", - "service_tier_flex", - "service_tier_priority", - "web_search", - "stream", - "stream_no_usage", - "response_model_override", - "stream_response_model_override", - "tool_call", - "stream_no_usage_tool_call", - "stream_no_usage_image_input", - "stream_no_usage_incomplete", - "stream_unvalidated", - "stream_no_usage_unvalidated", - "prompt_blocked", - "stream_prompt_blocked", -] - - -@dataclass(frozen=True, slots=True) -class Case: - name: CaseName - usage: ScriptedUsage - stream: bool = False - stream_usage: Literal["final_chunk", "absent"] = "final_chunk" - service_tier: Literal["flex", "priority"] | None = None - # For web_search the wire's reported call count is not always what gets - # billed: chat-completions surfaces only expose url_citation annotations, so - # the biller floors to one call; responses/messages/gemini report a real - # count. - billed_web_search_calls: int = 0 - response_model_override: bool = False - exact_spend: bool = True - tool_call: bool = False - image_input: bool = False - terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" - - def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: - return Scenario( - scenario_id=scenario_id, - wire=model.wire, - usage=self.usage, - model=model.provider_model, - output=ScriptedOutput( - text=text, - response_model=model.override_model if self.response_model_override else None, - tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) - if self.tool_call - else None, - terminal=self.terminal, - ), - stream_usage=self.stream_usage, - service_tier=self.service_tier, - ) - - -_BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) - TOOL_CALL_ARGUMENTS: Final = json.dumps({ "city": "Berlin", "days": 7, @@ -422,284 +375,9 @@ TOOL_CALL_ARGUMENTS: Final = json.dumps({ "notes": "filler " * 30, }) -_PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tokens=0) - - -def _web_search_case(model: FrontierModel) -> Case: - counts_exactly: Final = model.wire in ( - "openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate" - ) - return Case( - name="web_search", - usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), - billed_web_search_calls=3 if counts_exactly else 1, - ) - def cases_for(model: FrontierModel) -> tuple[Case, ...]: - rates: Final = model.rates - caps: Final = _WIRE_CAPS[model.wire] - candidates: Final[tuple[Case | None, ...]] = ( - Case(name="basic", usage=_BASIC_USAGE), - ( - Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) - if rates.cache_read_input_token_cost is not None and "cache_read" in caps - else None - ), - ( - Case( - name="cache_write_5m", - usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), - ) - if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps - else None - ), - ( - Case( - name="cache_write_1h", - usage=ScriptedUsage( - fresh_input_tokens=90, - cache_write_5m_tokens=20, - cache_write_1h_tokens=40, - output_tokens=30, - ), - ) - if ( - rates.cache_creation_input_token_cost_above_1hr is not None - and rates.cache_creation_input_token_cost is not None - and "cache_write_1h" in caps - ) - else None - ), - ( - Case( - name="reasoning", - usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), - ) - if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps - else None - ), - ( - Case( - name="audio", - usage=ScriptedUsage( - fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 - ), - ) - if ( - rates.input_cost_per_audio_token is not None - and rates.output_cost_per_audio_token is not None - and "audio" in caps - ) - else None - ), - ( - Case( - name="tiered", - usage=ScriptedUsage( - fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 - ), - ) - if ( - rates.input_cost_per_token_above_200k_tokens is not None - and rates.output_cost_per_token_above_200k_tokens is not None - ) - else None - ), - ( - Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") - if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None - else None - ), - ( - Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") - if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None - else None - ), - _web_search_case(model) if rates.search_context_cost_per_query is not None and "web_search" in caps else None, - Case(name="stream", usage=_BASIC_USAGE, stream=True), - ( - Case( - name="stream_no_usage", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - exact_spend=False, - ) - if "absent_usage" in caps - else None - ), - ( - Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True) - if "response_model" in caps - else None - ), - ( - Case( - name="stream_response_model_override", - usage=_BASIC_USAGE, - stream=True, - response_model_override=True, - ) - if "response_model" in caps - else None - ), - ( - Case(name="tool_call", usage=_BASIC_USAGE, tool_call=True) - if "tool_call" in caps - else None - ), - ( - Case( - name="stream_no_usage_tool_call", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - tool_call=True, - exact_spend=False, - ) - if "absent_usage" in caps and "tool_call" in caps - else None - ), - ( - Case( - name="stream_no_usage_image_input", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - image_input=True, - exact_spend=False, - ) - if "absent_usage" in caps and "image_input" in caps - else None - ), - ( - Case( - name="stream_no_usage_incomplete", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - terminal="incomplete", - exact_spend=False, - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="stream_unvalidated", - usage=_BASIC_USAGE, - stream=True, - terminal="unvalidated", - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="stream_no_usage_unvalidated", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - terminal="unvalidated", - exact_spend=False, - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="prompt_blocked", - usage=_PROMPT_BLOCKED_USAGE, - terminal="prompt_blocked", - response_model_override=True, - ) - if "prompt_blocked" in caps - else None - ), - ( - Case( - name="stream_prompt_blocked", - usage=_PROMPT_BLOCKED_USAGE, - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ) - if "prompt_blocked" in caps - else None - ), - ) - return tuple(case for case in candidates if case is not None) - - -@dataclass(frozen=True, slots=True) -class ExpectedCost: - """The expected bill split the way the spend row's cost_breakdown reports - it: the gross input component (cache reads/writes folded in), the output - component, and the tool-usage component.""" - - input_cost: float - output_cost: float - tool_cost: float - - @property - def total(self) -> float: - return self.input_cost + self.output_cost + self.tool_cost - - -def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: - """Literal arithmetic on the test-map rates over the scripted token counts. - - Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; - output = text*out + reasoning*reasoning + audio_out*audio_out; plus the - billed web-search calls at the medium search-context rate. Above-threshold - swaps every input/output rate to its ``_above_200k_tokens`` variant when - total prompt tokens exceed the threshold; a service tier swaps input/output - to the tier's variants, falling back to the base rate when a variant is - unset -- mirroring _get_token_base_cost in litellm's cost calculator. - """ - rates: Final = model.override_rates if case.response_model_override else model.rates - u: Final = case.usage - prompt_tokens: Final = ( - u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens - + u.cache_write_1h_tokens + u.audio_input_tokens - ) - tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS - in_rate: Final = ( - (rates.input_cost_per_token_above_200k_tokens if tiered else None) - or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.input_cost_per_token - or 0.0 - ) - out_rate: Final = ( - (rates.output_cost_per_token_above_200k_tokens if tiered else None) - or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.output_cost_per_token - or 0.0 - ) - input_cost: Final = ( - u.fresh_input_tokens * in_rate - + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) - + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) - ) - output_cost: Final = ( - u.output_tokens * out_rate - + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) - + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) - ) - search: Final = rates.search_context_cost_per_query - tool_cost: Final = case.billed_web_search_calls * ( - search.search_context_size_medium if search and search.search_context_size_medium else 0.0 - ) - return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) - - -def expected_cost(model: FrontierModel, case: Case) -> float: - return expected_breakdown(model, case).total + return tuple(case for case in CASES if case.applies_to(model)) def recount_cost( @@ -738,31 +416,23 @@ def image_input_data_url() -> str: IMAGE_INPUT_DATA_URL: Final = image_input_data_url() -def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: - """(prompt_tokens, completion_tokens) the spend row should carry, per the - wire's normalization: Anthropic folds cache read/write into prompt_tokens, - everyone else reports the totals the wire emitted.""" - u: Final = case.usage - if model.wire in ("anthropic_messages", "bedrock_converse"): - return ( - u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, - u.output_tokens, - ) - if model.wire in ("gemini_generate", "vertex_generate"): - return ( - u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, - u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, - ) - if model.wire == "openai_responses": - return ( - u.fresh_input_tokens + u.cache_read_tokens, - u.output_tokens + u.reasoning_tokens, - ) - return ( - u.fresh_input_tokens - + u.cache_read_tokens - + u.cache_write_5m_tokens - + u.cache_write_1h_tokens - + u.audio_input_tokens, - u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, - ) +class _ExpectedCell(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + input_cost: float + output_cost: float + prompt_tokens: int + completion_tokens: int + + +_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, _ExpectedCell]) +EXPECTED: Final[Mapping[str, _ExpectedCell]] = MappingProxyType( + _EXPECTED_ADAPTER.validate_python(json.loads(EXPECTED_PATH.read_text())) + if EXPECTED_PATH.exists() + else {} +) + + +def expected_key(model: FrontierModel, case: Case) -> str: + return f"{model.map_key}|{case.name}" diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json new file mode 100644 index 00000000000..7a92fb2476f --- /dev/null +++ b/tests/e2e/cost_calculation/expected.json @@ -0,0 +1,2004 @@ +{ + "anthropic.claude-sonnet-5-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.03128, + "output_cost": 0.0085, + "prompt_tokens": 150, + "spend": 0.03978 + }, + "anthropic.claude-sonnet-5-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "anthropic.claude-sonnet-5-v1:0|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01785, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.028050000000000002 + }, + "anthropic.claude-sonnet-5-v1:0|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.052700000000000004, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.06290000000000001 + }, + "anthropic.claude-sonnet-5-v1:0|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0459, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.056100000000000004 + }, + "anthropic.claude-sonnet-5-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "anthropic.claude-sonnet-5-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.013600000000000001, + "output_cost": 0.0085, + "prompt_tokens": 80, + "spend": 0.0221 + }, + "anthropic.claude-sonnet-5-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "azure/gpt-5.4-mini|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.03424, + "output_cost": 0.02336, + "prompt_tokens": 155, + "spend": 0.0576 + }, + "azure/gpt-5.4-mini|audio": { + "completion_tokens": 45, + "input_cost": 0.04, + "output_cost": 0.0264, + "prompt_tokens": 125, + "spend": 0.0664 + }, + "azure/gpt-5.4-mini|basic": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0168, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0264 + }, + "azure/gpt-5.4-mini|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.049600000000000005, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0592 + }, + "azure/gpt-5.4-mini|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0432, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0528 + }, + "azure/gpt-5.4-mini|reasoning": { + "completion_tokens": 100, + "input_cost": 0.016, + "output_cost": 0.0656, + "prompt_tokens": 100, + "spend": 0.0816 + }, + "azure/gpt-5.4-mini|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0288, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.0448 + }, + "azure/gpt-5.4-mini|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.03264, + "output_cost": 0.01728, + "prompt_tokens": 120, + "spend": 0.049920000000000006 + }, + "azure/gpt-5.4-mini|stream": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0128, + "output_cost": 0.008, + "prompt_tokens": 80, + "spend": 0.0208 + }, + "azure/gpt-5.4-mini|tiered": { + "completion_tokens": 30, + "input_cost": 256.00128, + "output_cost": 0.0432, + "prompt_tokens": 200001, + "spend": 256.04448 + }, + "azure/gpt-5.4-mini|tool_call": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|web_search": { + "completion_tokens": 30, + "input_cost": 0.016, + "output_cost": 0.009600000000000001, + "prompt_tokens": 100, + "spend": 0.0456 + }, + "azure/gpt-5.6|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.0321, + "output_cost": 0.0219, + "prompt_tokens": 155, + "spend": 0.05399999999999999 + }, + "azure/gpt-5.6|audio": { + "completion_tokens": 45, + "input_cost": 0.0375, + "output_cost": 0.02475, + "prompt_tokens": 125, + "spend": 0.06225 + }, + "azure/gpt-5.6|basic": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01575, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.02475 + }, + "azure/gpt-5.6|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0465, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.0555 + }, + "azure/gpt-5.6|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.040499999999999994, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.049499999999999995 + }, + "azure/gpt-5.6|reasoning": { + "completion_tokens": 100, + "input_cost": 0.015, + "output_cost": 0.0615, + "prompt_tokens": 100, + "spend": 0.0765 + }, + "azure/gpt-5.6|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.6|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.027, + "output_cost": 0.015, + "prompt_tokens": 120, + "spend": 0.041999999999999996 + }, + "azure/gpt-5.6|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.030600000000000002, + "output_cost": 0.0162, + "prompt_tokens": 120, + "spend": 0.0468 + }, + "azure/gpt-5.6|stream": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.6|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.011999999999999999, + "output_cost": 0.0075, + "prompt_tokens": 80, + "spend": 0.019499999999999997 + }, + "azure/gpt-5.6|tiered": { + "completion_tokens": 30, + "input_cost": 240.00119999999998, + "output_cost": 0.0405, + "prompt_tokens": 200001, + "spend": 240.0417 + }, + "azure/gpt-5.6|tool_call": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|web_search": { + "completion_tokens": 30, + "input_cost": 0.015, + "output_cost": 0.009, + "prompt_tokens": 100, + "spend": 0.044 + }, + "claude-haiku-4-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.012880000000000003, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 150, + "spend": 0.016380000000000002 + }, + "claude-haiku-4-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.012880000000000003, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 150, + "spend": 0.016380000000000002 + }, + "claude-haiku-4-5|basic": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.007350000000000001, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.011550000000000001 + }, + "claude-haiku-4-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.021700000000000004, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.025900000000000006 + }, + "claude-haiku-4-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0189, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.023100000000000002 + }, + "claude-haiku-4-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-haiku-4-5|stream": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-haiku-4-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.005600000000000001, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 80, + "spend": 0.0091 + }, + "claude-haiku-4-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.007000000000000001, + "output_cost": 0.004200000000000001, + "prompt_tokens": 100, + "spend": 0.0712 + }, + "claude-opus-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.0092, + "output_cost": 0.0025, + "prompt_tokens": 150, + "spend": 0.0117 + }, + "claude-opus-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.0092, + "output_cost": 0.0025, + "prompt_tokens": 150, + "spend": 0.0117 + }, + "claude-opus-5|basic": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.00525, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.00825 + }, + "claude-opus-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0155, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.0185 + }, + "claude-opus-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.013500000000000002, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.0165 + }, + "claude-opus-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-opus-5|stream": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-opus-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.004, + "output_cost": 0.0025, + "prompt_tokens": 80, + "spend": 0.006500000000000001 + }, + "claude-opus-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.005, + "output_cost": 0.003, + "prompt_tokens": 100, + "spend": 0.068 + }, + "claude-sonnet-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.011040000000000001, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 150, + "spend": 0.014040000000000002 + }, + "claude-sonnet-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.011040000000000001, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 150, + "spend": 0.014040000000000002 + }, + "claude-sonnet-5|basic": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.006300000000000001, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0099 + }, + "claude-sonnet-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.018600000000000002, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0222 + }, + "claude-sonnet-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.016200000000000003, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0198 + }, + "claude-sonnet-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-sonnet-5|stream": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-sonnet-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 80, + "spend": 0.007800000000000001 + }, + "claude-sonnet-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.006000000000000001, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 100, + "spend": 0.0696 + }, + "fireworks_ai/deepseek-v4p1-flash|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.011760000000000001, + "output_cost": 0.007000000000000001, + "prompt_tokens": 120, + "spend": 0.018760000000000002 + }, + "fireworks_ai/deepseek-v4p1-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.030500000000000003, + "output_cost": 0.019950000000000002, + "prompt_tokens": 125, + "spend": 0.05045000000000001 + }, + "fireworks_ai/deepseek-v4p1-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.014700000000000001, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.023100000000000002 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0368, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.045200000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0324, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.0408 + }, + "fireworks_ai/deepseek-v4p1-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.014000000000000002, + "output_cost": 0.0469, + "prompt_tokens": 100, + "spend": 0.060899999999999996 + }, + "fireworks_ai/deepseek-v4p1-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/deepseek-v4p1-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/deepseek-v4p1-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.011200000000000002, + "output_cost": 0.007000000000000001, + "prompt_tokens": 80, + "spend": 0.0182 + }, + "fireworks_ai/deepseek-v4p1-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.014000000000000002, + "output_cost": 0.008400000000000001, + "prompt_tokens": 100, + "spend": 0.04240000000000001 + }, + "fireworks_ai/kimi-k3|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.01008, + "output_cost": 0.006000000000000001, + "prompt_tokens": 120, + "spend": 0.01608 + }, + "fireworks_ai/kimi-k3|audio": { + "completion_tokens": 45, + "input_cost": 0.028500000000000004, + "output_cost": 0.01875, + "prompt_tokens": 125, + "spend": 0.04725 + }, + "fireworks_ai/kimi-k3|basic": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.012600000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0198 + }, + "fireworks_ai/kimi-k3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.035, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0422 + }, + "fireworks_ai/kimi-k3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.030600000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0378 + }, + "fireworks_ai/kimi-k3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.012000000000000002, + "output_cost": 0.0457, + "prompt_tokens": 100, + "spend": 0.0577 + }, + "fireworks_ai/kimi-k3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/kimi-k3|stream": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/kimi-k3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.009600000000000001, + "output_cost": 0.006000000000000001, + "prompt_tokens": 80, + "spend": 0.015600000000000003 + }, + "fireworks_ai/kimi-k3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|web_search": { + "completion_tokens": 30, + "input_cost": 0.012000000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 100, + "spend": 0.0392 + }, + "fireworks_ai/qwen3p8-max|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.010920000000000001, + "output_cost": 0.006500000000000001, + "prompt_tokens": 120, + "spend": 0.01742 + }, + "fireworks_ai/qwen3p8-max|audio": { + "completion_tokens": 45, + "input_cost": 0.029500000000000002, + "output_cost": 0.01935, + "prompt_tokens": 125, + "spend": 0.048850000000000005 + }, + "fireworks_ai/qwen3p8-max|basic": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01365, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.021450000000000004 + }, + "fireworks_ai/qwen3p8-max|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0359, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.0437 + }, + "fireworks_ai/qwen3p8-max|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0315, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.0393 + }, + "fireworks_ai/qwen3p8-max|reasoning": { + "completion_tokens": 100, + "input_cost": 0.013000000000000001, + "output_cost": 0.0463, + "prompt_tokens": 100, + "spend": 0.059300000000000005 + }, + "fireworks_ai/qwen3p8-max|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/qwen3p8-max|stream": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/qwen3p8-max|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.010400000000000001, + "output_cost": 0.006500000000000001, + "prompt_tokens": 80, + "spend": 0.016900000000000002 + }, + "fireworks_ai/qwen3p8-max|tool_call": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|web_search": { + "completion_tokens": 30, + "input_cost": 0.013000000000000001, + "output_cost": 0.007800000000000001, + "prompt_tokens": 100, + "spend": 0.0408 + }, + "gemini-3.1-pro-preview|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.023940000000000003, + "output_cost": 0.030660000000000003, + "prompt_tokens": 125, + "spend": 0.05460000000000001 + }, + "gemini-3.1-pro-preview|audio": { + "completion_tokens": 45, + "input_cost": 0.052500000000000005, + "output_cost": 0.03465, + "prompt_tokens": 125, + "spend": 0.08715 + }, + "gemini-3.1-pro-preview|basic": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|cache_read": { + "completion_tokens": 30, + "input_cost": 0.02205, + "output_cost": 0.0126, + "prompt_tokens": 150, + "spend": 0.03465 + }, + "gemini-3.1-pro-preview|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.2, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.2 + }, + "gemini-3.1-pro-preview|reasoning": { + "completion_tokens": 100, + "input_cost": 0.021, + "output_cost": 0.0861, + "prompt_tokens": 100, + "spend": 0.1071 + }, + "gemini-3.1-pro-preview|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.1-pro-preview|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0378, + "output_cost": 0.020999999999999998, + "prompt_tokens": 120, + "spend": 0.0588 + }, + "gemini-3.1-pro-preview|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.04284, + "output_cost": 0.02268, + "prompt_tokens": 120, + "spend": 0.06552 + }, + "gemini-3.1-pro-preview|stream": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.2, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.2 + }, + "gemini-3.1-pro-preview|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.1-pro-preview|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.016800000000000002, + "output_cost": 0.0105, + "prompt_tokens": 80, + "spend": 0.027300000000000005 + }, + "gemini-3.1-pro-preview|tiered": { + "completion_tokens": 30, + "input_cost": 336.00168, + "output_cost": 0.0567, + "prompt_tokens": 200001, + "spend": 336.05838 + }, + "gemini-3.1-pro-preview|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|web_search": { + "completion_tokens": 30, + "input_cost": 0.021, + "output_cost": 0.0126, + "prompt_tokens": 100, + "spend": 0.0936 + }, + "gemini-3.8-flash|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.022799999999999997, + "output_cost": 0.0292, + "prompt_tokens": 125, + "spend": 0.052 + }, + "gemini-3.8-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.05, + "output_cost": 0.033, + "prompt_tokens": 125, + "spend": 0.083 + }, + "gemini-3.8-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.021, + "output_cost": 0.012, + "prompt_tokens": 150, + "spend": 0.033 + }, + "gemini-3.8-flash|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.21000000000000002, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.21000000000000002 + }, + "gemini-3.8-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.02, + "output_cost": 0.082, + "prompt_tokens": 100, + "spend": 0.10200000000000001 + }, + "gemini-3.8-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.8-flash|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.036, + "output_cost": 0.02, + "prompt_tokens": 120, + "spend": 0.055999999999999994 + }, + "gemini-3.8-flash|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.0408, + "output_cost": 0.0216, + "prompt_tokens": 120, + "spend": 0.062400000000000004 + }, + "gemini-3.8-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.21000000000000002, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.21000000000000002 + }, + "gemini-3.8-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.8-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.016, + "output_cost": 0.01, + "prompt_tokens": 80, + "spend": 0.026000000000000002 + }, + "gemini-3.8-flash|tiered": { + "completion_tokens": 30, + "input_cost": 320.0016, + "output_cost": 0.054, + "prompt_tokens": 200001, + "spend": 320.05559999999997 + }, + "gemini-3.8-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.02, + "output_cost": 0.012, + "prompt_tokens": 100, + "spend": 0.092 + }, + "gemini/gemini-3.1-pro-preview|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.010260000000000002, + "output_cost": 0.01314, + "prompt_tokens": 125, + "spend": 0.023400000000000004 + }, + "gemini/gemini-3.1-pro-preview|audio": { + "completion_tokens": 45, + "input_cost": 0.0225, + "output_cost": 0.014849999999999999, + "prompt_tokens": 125, + "spend": 0.037349999999999994 + }, + "gemini/gemini-3.1-pro-preview|basic": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|cache_read": { + "completion_tokens": 30, + "input_cost": 0.009450000000000002, + "output_cost": 0.0054, + "prompt_tokens": 150, + "spend": 0.014850000000000002 + }, + "gemini/gemini-3.1-pro-preview|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.08, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.08 + }, + "gemini/gemini-3.1-pro-preview|reasoning": { + "completion_tokens": 100, + "input_cost": 0.009000000000000001, + "output_cost": 0.0369, + "prompt_tokens": 100, + "spend": 0.0459 + }, + "gemini/gemini-3.1-pro-preview|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.1-pro-preview|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0162, + "output_cost": 0.009000000000000001, + "prompt_tokens": 120, + "spend": 0.0252 + }, + "gemini/gemini-3.1-pro-preview|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.01836, + "output_cost": 0.00972, + "prompt_tokens": 120, + "spend": 0.02808 + }, + "gemini/gemini-3.1-pro-preview|stream": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.08, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.08 + }, + "gemini/gemini-3.1-pro-preview|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.1-pro-preview|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.007200000000000001, + "output_cost": 0.0045000000000000005, + "prompt_tokens": 80, + "spend": 0.011700000000000002 + }, + "gemini/gemini-3.1-pro-preview|tiered": { + "completion_tokens": 30, + "input_cost": 144.00072, + "output_cost": 0.024300000000000002, + "prompt_tokens": 200001, + "spend": 144.02502 + }, + "gemini/gemini-3.1-pro-preview|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|web_search": { + "completion_tokens": 30, + "input_cost": 0.009000000000000001, + "output_cost": 0.0054, + "prompt_tokens": 100, + "spend": 0.0744 + }, + "gemini/gemini-3.8-flash|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.00912, + "output_cost": 0.01168, + "prompt_tokens": 125, + "spend": 0.0208 + }, + "gemini/gemini-3.8-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.02, + "output_cost": 0.0132, + "prompt_tokens": 125, + "spend": 0.0332 + }, + "gemini/gemini-3.8-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0084, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 150, + "spend": 0.0132 + }, + "gemini/gemini-3.8-flash|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.09000000000000001, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.09000000000000001 + }, + "gemini/gemini-3.8-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.008, + "output_cost": 0.0328, + "prompt_tokens": 100, + "spend": 0.0408 + }, + "gemini/gemini-3.8-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.8-flash|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0144, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.0224 + }, + "gemini/gemini-3.8-flash|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.01632, + "output_cost": 0.00864, + "prompt_tokens": 120, + "spend": 0.024960000000000003 + }, + "gemini/gemini-3.8-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.09000000000000001, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.09000000000000001 + }, + "gemini/gemini-3.8-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.8-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0064, + "output_cost": 0.004, + "prompt_tokens": 80, + "spend": 0.0104 + }, + "gemini/gemini-3.8-flash|tiered": { + "completion_tokens": 30, + "input_cost": 128.00064, + "output_cost": 0.0216, + "prompt_tokens": 200001, + "spend": 128.02224 + }, + "gemini/gemini-3.8-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.008, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 100, + "spend": 0.0728 + }, + "gpt-5.3-codex|all_components_responses": { + "completion_tokens": 40, + "input_cost": 0.00252, + "output_cost": 0.0037500000000000007, + "prompt_tokens": 120, + "spend": 0.006270000000000001 + }, + "gpt-5.3-codex|basic": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0031500000000000005, + "output_cost": 0.0018000000000000002, + "prompt_tokens": 150, + "spend": 0.00495 + }, + "gpt-5.3-codex|reasoning": { + "completion_tokens": 100, + "input_cost": 0.0030000000000000005, + "output_cost": 0.0123, + "prompt_tokens": 100, + "spend": 0.015300000000000001 + }, + "gpt-5.3-codex|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.3-codex|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0054, + "output_cost": 0.003, + "prompt_tokens": 120, + "spend": 0.008400000000000001 + }, + "gpt-5.3-codex|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00612, + "output_cost": 0.00324, + "prompt_tokens": 120, + "spend": 0.00936 + }, + "gpt-5.3-codex|stream": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|stream_incomplete": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.3-codex|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0015000000000000002, + "prompt_tokens": 80, + "spend": 0.0039000000000000007 + }, + "gpt-5.3-codex|stream_unvalidated": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|tiered": { + "completion_tokens": 30, + "input_cost": 48.000240000000005, + "output_cost": 0.0081, + "prompt_tokens": 200001, + "spend": 48.008340000000004 + }, + "gpt-5.3-codex|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|web_search": { + "completion_tokens": 30, + "input_cost": 0.0030000000000000005, + "output_cost": 0.0018000000000000002, + "prompt_tokens": 100, + "spend": 0.0648 + }, + "gpt-5.4-mini|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.00856, + "output_cost": 0.00584, + "prompt_tokens": 155, + "spend": 0.0144 + }, + "gpt-5.4-mini|audio": { + "completion_tokens": 45, + "input_cost": 0.01, + "output_cost": 0.0066, + "prompt_tokens": 125, + "spend": 0.0166 + }, + "gpt-5.4-mini|basic": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0042, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0066 + }, + "gpt-5.4-mini|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.012400000000000001, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0148 + }, + "gpt-5.4-mini|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0108, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0132 + }, + "gpt-5.4-mini|reasoning": { + "completion_tokens": 100, + "input_cost": 0.004, + "output_cost": 0.0164, + "prompt_tokens": 100, + "spend": 0.0204 + }, + "gpt-5.4-mini|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.4-mini|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0072, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.0112 + }, + "gpt-5.4-mini|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00816, + "output_cost": 0.00432, + "prompt_tokens": 120, + "spend": 0.012480000000000002 + }, + "gpt-5.4-mini|stream": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.4-mini|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0032, + "output_cost": 0.002, + "prompt_tokens": 80, + "spend": 0.0052 + }, + "gpt-5.4-mini|tiered": { + "completion_tokens": 30, + "input_cost": 64.00032, + "output_cost": 0.0108, + "prompt_tokens": 200001, + "spend": 64.01112 + }, + "gpt-5.4-mini|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|web_search": { + "completion_tokens": 30, + "input_cost": 0.004, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 100, + "spend": 0.0264 + }, + "gpt-5.5-pro|all_components_responses": { + "completion_tokens": 40, + "input_cost": 0.00168, + "output_cost": 0.0025, + "prompt_tokens": 120, + "spend": 0.00418 + }, + "gpt-5.5-pro|basic": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0021, + "output_cost": 0.0012000000000000001, + "prompt_tokens": 150, + "spend": 0.0033 + }, + "gpt-5.5-pro|reasoning": { + "completion_tokens": 100, + "input_cost": 0.002, + "output_cost": 0.0082, + "prompt_tokens": 100, + "spend": 0.0102 + }, + "gpt-5.5-pro|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.5-pro|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0036, + "output_cost": 0.002, + "prompt_tokens": 120, + "spend": 0.0056 + }, + "gpt-5.5-pro|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00408, + "output_cost": 0.00216, + "prompt_tokens": 120, + "spend": 0.006240000000000001 + }, + "gpt-5.5-pro|stream": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|stream_incomplete": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.5-pro|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0016, + "output_cost": 0.001, + "prompt_tokens": 80, + "spend": 0.0026 + }, + "gpt-5.5-pro|stream_unvalidated": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|tiered": { + "completion_tokens": 30, + "input_cost": 32.00016, + "output_cost": 0.0054, + "prompt_tokens": 200001, + "spend": 32.00556 + }, + "gpt-5.5-pro|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|web_search": { + "completion_tokens": 30, + "input_cost": 0.002, + "output_cost": 0.0012000000000000001, + "prompt_tokens": 100, + "spend": 0.06319999999999999 + }, + "gpt-5.6|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.00214, + "output_cost": 0.00146, + "prompt_tokens": 155, + "spend": 0.0036 + }, + "gpt-5.6|audio": { + "completion_tokens": 45, + "input_cost": 0.0025, + "output_cost": 0.00165, + "prompt_tokens": 125, + "spend": 0.00415 + }, + "gpt-5.6|basic": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|cache_read": { + "completion_tokens": 30, + "input_cost": 0.00105, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.00165 + }, + "gpt-5.6|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0031000000000000003, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.0037 + }, + "gpt-5.6|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0027, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.0033 + }, + "gpt-5.6|reasoning": { + "completion_tokens": 100, + "input_cost": 0.001, + "output_cost": 0.0041, + "prompt_tokens": 100, + "spend": 0.0051 + }, + "gpt-5.6|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.6|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0018, + "output_cost": 0.001, + "prompt_tokens": 120, + "spend": 0.0028 + }, + "gpt-5.6|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00204, + "output_cost": 0.00108, + "prompt_tokens": 120, + "spend": 0.0031200000000000004 + }, + "gpt-5.6|stream": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.6|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0008, + "output_cost": 0.0005, + "prompt_tokens": 80, + "spend": 0.0013 + }, + "gpt-5.6|tiered": { + "completion_tokens": 30, + "input_cost": 16.00008, + "output_cost": 0.0027, + "prompt_tokens": 200001, + "spend": 16.00278 + }, + "gpt-5.6|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|web_search": { + "completion_tokens": 30, + "input_cost": 0.001, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 100, + "spend": 0.0216 + }, + "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.020900000000000002, + "output_cost": 0.0095, + "prompt_tokens": 150, + "spend": 0.030400000000000003 + }, + "meta.llama4-maverick-17b-instruct-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "meta.llama4-maverick-17b-instruct-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "meta.llama4-maverick-17b-instruct-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.015200000000000002, + "output_cost": 0.0095, + "prompt_tokens": 80, + "spend": 0.0247 + }, + "meta.llama4-maverick-17b-instruct-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "together_ai/moonshotai/Kimi-K3|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.0214, + "output_cost": 0.0146, + "prompt_tokens": 155, + "spend": 0.036 + }, + "together_ai/moonshotai/Kimi-K3|audio": { + "completion_tokens": 45, + "input_cost": 0.025, + "output_cost": 0.0165, + "prompt_tokens": 125, + "spend": 0.0415 + }, + "together_ai/moonshotai/Kimi-K3|basic": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0105, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.0165 + }, + "together_ai/moonshotai/Kimi-K3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.031, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.037 + }, + "together_ai/moonshotai/Kimi-K3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.027000000000000003, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.033 + }, + "together_ai/moonshotai/Kimi-K3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.01, + "output_cost": 0.041, + "prompt_tokens": 100, + "spend": 0.051000000000000004 + }, + "together_ai/moonshotai/Kimi-K3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/moonshotai/Kimi-K3|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.018000000000000002, + "output_cost": 0.01, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "together_ai/moonshotai/Kimi-K3|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.0108, + "prompt_tokens": 120, + "spend": 0.031200000000000002 + }, + "together_ai/moonshotai/Kimi-K3|stream": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/moonshotai/Kimi-K3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.008, + "output_cost": 0.005, + "prompt_tokens": 80, + "spend": 0.013000000000000001 + }, + "together_ai/moonshotai/Kimi-K3|tiered": { + "completion_tokens": 30, + "input_cost": 160.0008, + "output_cost": 0.027000000000000003, + "prompt_tokens": 200001, + "spend": 160.02779999999998 + }, + "together_ai/moonshotai/Kimi-K3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|web_search": { + "completion_tokens": 30, + "input_cost": 0.01, + "output_cost": 0.006, + "prompt_tokens": 100, + "spend": 0.036000000000000004 + }, + "together_ai/zai-org/GLM-5.3|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.023540000000000002, + "output_cost": 0.01606, + "prompt_tokens": 155, + "spend": 0.0396 + }, + "together_ai/zai-org/GLM-5.3|audio": { + "completion_tokens": 45, + "input_cost": 0.027500000000000004, + "output_cost": 0.01815, + "prompt_tokens": 125, + "spend": 0.04565 + }, + "together_ai/zai-org/GLM-5.3|basic": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.011550000000000001, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.01815 + }, + "together_ai/zai-org/GLM-5.3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.034100000000000005, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.04070000000000001 + }, + "together_ai/zai-org/GLM-5.3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.029699999999999997, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.0363 + }, + "together_ai/zai-org/GLM-5.3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.011000000000000001, + "output_cost": 0.0451, + "prompt_tokens": 100, + "spend": 0.056100000000000004 + }, + "together_ai/zai-org/GLM-5.3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/zai-org/GLM-5.3|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.019799999999999998, + "output_cost": 0.011000000000000001, + "prompt_tokens": 120, + "spend": 0.0308 + }, + "together_ai/zai-org/GLM-5.3|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.022439999999999998, + "output_cost": 0.01188, + "prompt_tokens": 120, + "spend": 0.034319999999999996 + }, + "together_ai/zai-org/GLM-5.3|stream": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/zai-org/GLM-5.3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0088, + "output_cost": 0.0055000000000000005, + "prompt_tokens": 80, + "spend": 0.0143 + }, + "together_ai/zai-org/GLM-5.3|tiered": { + "completion_tokens": 30, + "input_cost": 176.00088, + "output_cost": 0.0297, + "prompt_tokens": 200001, + "spend": 176.03058 + }, + "together_ai/zai-org/GLM-5.3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|web_search": { + "completion_tokens": 30, + "input_cost": 0.011000000000000001, + "output_cost": 0.0066, + "prompt_tokens": 100, + "spend": 0.0376 + }, + "us.anthropic.claude-opus-5-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.033120000000000004, + "output_cost": 0.009000000000000001, + "prompt_tokens": 150, + "spend": 0.042120000000000005 + }, + "us.anthropic.claude-opus-5-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|cache_read": { + "completion_tokens": 30, + "input_cost": 0.018900000000000004, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.029700000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0558, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.0666 + }, + "us.anthropic.claude-opus-5-v1:0|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.048600000000000004, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.05940000000000001 + }, + "us.anthropic.claude-opus-5-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.014400000000000001, + "output_cost": 0.009000000000000001, + "prompt_tokens": 80, + "spend": 0.023400000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + } +} diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py new file mode 100644 index 00000000000..de979f272fe --- /dev/null +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -0,0 +1,189 @@ +"""Golden generator for the cost suite. Run: + + uv run python tests/e2e/cost_calculation/generate_expected.py + +Loads the derived matrix (models x applicable cases), computes the golden for +each exact-spend cell from the rate arithmetic, and writes ``expected.json`` +with sorted keys. Default behaviour adds missing cells and drops stale cells +but never overwrites an existing cell's values (a reviewed golden is +authoritative); ``--rewrite`` recomputes everything. Prints added/removed/kept +counts. +""" + +from __future__ import annotations + +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from cost_matrix import ( # noqa: E402 # path bootstrap before package-local imports + EXPECTED_PATH, + FRONTIER_MODELS, + TIER_THRESHOLD_TOKENS, + Case, + CostMapEntry, + FrontierModel, + cases_for, + expected_key, +) + +# Wires whose response surface reports a real web-search call count; the +# chat-completions wires only expose url_citation annotations, so their billed +# count floors to one. +_EXACT_WEB_SEARCH_WIRES: Final = frozenset( + {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"} +) + + +def billed_web_search_calls(model: FrontierModel, case: Case) -> int: + if case.usage.web_search_calls == 0: + return 0 + return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1 + + +@dataclass(frozen=True, slots=True) +class ExpectedCost: + """The expected bill split the way the spend row's cost_breakdown reports + it: the gross input component (cache reads/writes folded in), the output + component, and the tool-usage component.""" + + input_cost: float + output_cost: float + tool_cost: float + + @property + def total(self) -> float: + return self.input_cost + self.output_cost + self.tool_cost + + +def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: + """Literal arithmetic on the test-map rates over the scripted token counts. + + Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; + output = text*out + reasoning*reasoning + audio_out*audio_out; plus the + billed web-search calls at the medium search-context rate. Above-threshold + swaps every input/output rate to its ``_above_200k_tokens`` variant when + total prompt tokens exceed the threshold; a service tier swaps input/output + to the tier's variants, falling back to the base rate when a variant is + unset -- mirroring _get_token_base_cost in litellm's cost calculator. + """ + rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates + u: Final = case.usage + prompt_tokens: Final = ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + u.audio_input_tokens + ) + tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate: Final = ( + (rates.input_cost_per_token_above_200k_tokens if tiered else None) + or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.input_cost_per_token + or 0.0 + ) + out_rate: Final = ( + (rates.output_cost_per_token_above_200k_tokens if tiered else None) + or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.output_cost_per_token + or 0.0 + ) + # The biller charges cache writes at the input rate when the entry carries + # no cache_creation rate (cost_calculator.py:2452), and at the 5m write + # rate when the 1h variant is unset; cache reads bill only at their own + # rate (zero when the entry lacks one). + write_5m_rate: Final = rates.cache_creation_input_token_cost or in_rate + input_cost: Final = ( + u.fresh_input_tokens * in_rate + + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + + u.cache_write_5m_tokens * write_5m_rate + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or write_5m_rate) + + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + ) + output_cost: Final = ( + u.output_tokens * out_rate + + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + ) + search: Final = rates.search_context_cost_per_query + tool_cost: Final = billed_web_search_calls(model, case) * ( + search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + ) + return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) + + +def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: + """(prompt_tokens, completion_tokens) the spend row should carry, per the + wire's normalization: Anthropic folds cache read/write into prompt_tokens, + everyone else reports the totals the wire emitted.""" + u: Final = case.usage + if model.wire in ("anthropic_messages", "bedrock_converse"): + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, + u.output_tokens, + ) + if model.wire in ("gemini_generate", "vertex_generate"): + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + if model.wire == "openai_responses": + return ( + u.fresh_input_tokens + u.cache_read_tokens, + u.output_tokens + u.reasoning_tokens, + ) + return ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + + +def _proposed() -> dict[str, dict[str, object]]: + return { + expected_key(model, case): ( + lambda breakdown, tokens: { + "spend": breakdown.total, + "input_cost": breakdown.input_cost, + "output_cost": breakdown.output_cost, + "prompt_tokens": tokens[0], + "completion_tokens": tokens[1], + } + )(expected_breakdown(model, case), expected_token_columns(model, case)) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + + +def main() -> None: + rewrite: Final = "--rewrite" in sys.argv[1:] + proposed: Final = _proposed() + existing: Final = ( + json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} + ) + merged: Final = { + key: (proposed[key] if rewrite or key not in existing else existing[key]) + for key in sorted(proposed) + } + added: Final = sum(1 for key in proposed if key not in existing) + removed: Final = sum(1 for key in existing if key not in proposed) + kept: Final = sum(1 for key in proposed if key in existing and not rewrite) + rewritten: Final = sum(1 for key in proposed if key in existing and rewrite) + EXPECTED_PATH.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") + print( + f"expected.json: {added} added, {removed} removed, {kept} kept, " + f"{rewritten} rewritten ({len(merged)} cells)" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py new file mode 100644 index 00000000000..fdbb6ddd293 --- /dev/null +++ b/tests/e2e/cost_calculation/test_matrix_data.py @@ -0,0 +1,64 @@ +"""Freshness checks for the cost suite's data files; markerless, so it runs on +any pytest invocation of the folder without the stack. expected.json is the +oracle: these tests check its key set against the derived matrix, never its +values (the generator proposes, the file decides).""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from cost_matrix import ( + _CASES_FILE, + _COST_MAP, + CASES, + EXPECTED, + FRONTIER_MODELS, + CostMapEntry, + cases_for, + expected_key, +) + + +def test_expected_keys_match_derived_exact_cells() -> None: + derived: Final = { + expected_key(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + golden: Final = set(EXPECTED) + if derived != golden: + missing: Final = sorted(derived - golden) + stale: Final = sorted(golden - derived) + pytest.fail( + "expected.json is out of sync with the derived matrix; run " + "uv run python tests/e2e/cost_calculation/generate_expected.py " + f"(missing: {missing}; stale: {stale})" + ) + + +def test_deployments_reference_existing_map_keys() -> None: + unknown: Final = sorted( + spec.map_key for spec in _CASES_FILE.deployments if spec.map_key not in _COST_MAP + ) + assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" + + +def test_requires_rates_are_cost_map_fields() -> None: + fields: Final = set(CostMapEntry.model_fields) + unknown: Final = sorted( + {field for case in CASES for field in case.requires_rates} - fields + ) + assert not unknown, f"requires_rates names that are not CostMapEntry fields: {unknown}" + + +def test_no_two_entries_share_input_rate() -> None: + rates: Final = [ + entry.input_cost_per_token for entry in _COST_MAP.values() + ] + assert len(rates) == len(set(rates)), ( + "two cost_map entries share input_cost_per_token; the suite relies on " + "distinct rates so a wrong-model bill can never coincidentally match" + ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 0b4f3e1fd37..7cd128ad6fb 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -1,7 +1,7 @@ -"""Token-pricing e2e: every (frontier model, pricing-component case) cell runs a -scripted-usage call through a deployment registered on the cost-map proxy, and -the spend row plus response-cost header must equal literal arithmetic on the -test map's rates. +"""Token-pricing e2e: every (map entry, case) cell derived from cost_map.json x +cases.json runs a scripted-usage call through a deployment registered on the +cost-map proxy, and the spend row plus response-cost header must equal the +reviewed golden in expected.json verbatim -- no rate arithmetic lives here. Nothing here touches a real provider or the bundled cost map: the proxy's upstream is the scripted-provider sidecar and its entire cost map is @@ -15,13 +15,13 @@ from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( + EXPECTED, FRONTIER_MODELS, IMAGE_INPUT_DATA_URL, Case, FrontierModel, cases_for, - expected_cost, - expected_token_columns, + expected_key, recount_cost, ) from e2e_config import unique_marker @@ -110,17 +110,6 @@ class TestTokenPricing: ) assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected: Final = expected_cost(model, case) - if case.exact_spend and not case.stream: - # Streamed responses commit headers before the bill is computed, so - # the x-litellm-response-cost header is asserted only on non-stream - # calls. - assert response.response_cost is not None and cost_rows.approx_equal( - response.response_cost, expected - ), ( - f"x-litellm-response-cost {response.response_cost} != expected {expected}" - ) - row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, @@ -149,16 +138,39 @@ class TestTokenPricing: cost_rows.assert_total_is_sum_of_components(row) return - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( - f"{model.map_key}/{case.name}: spend {row.spend} != expected {expected} " + golden: Final = EXPECTED[expected_key(model, case)] + + if not case.stream: + # Streamed responses commit headers before the bill is computed, so + # the x-litellm-response-cost header is asserted only on non-stream + # calls. + assert response.response_cost is not None and cost_rows.approx_equal( + response.response_cost, golden.spend + ), ( + f"x-litellm-response-cost {response.response_cost} != golden {golden.spend}" + ) + + assert row.spend is not None and cost_rows.approx_equal(row.spend, golden.spend), ( + f"{model.map_key}/{case.name}: spend {row.spend} != golden {golden.spend} " f"(breakdown {row.breakdown.model_dump()})" ) - - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + breakdown: Final = row.breakdown + assert breakdown.input_cost is not None and cost_rows.approx_equal( + breakdown.input_cost, golden.input_cost + ), ( + f"{model.map_key}/{case.name}: gross input_cost {breakdown.input_cost} " + f"!= golden {golden.input_cost}; cached/written tokens billed at the input rate" ) - assert row.completion_tokens == completion_tokens, ( - f"completion_tokens {row.completion_tokens} != {completion_tokens}" + assert breakdown.output_cost is not None and cost_rows.approx_equal( + breakdown.output_cost, golden.output_cost + ), ( + f"{model.map_key}/{case.name}: output_cost {breakdown.output_cost} " + f"!= golden {golden.output_cost}" + ) + assert row.prompt_tokens == golden.prompt_tokens, ( + f"prompt_tokens {row.prompt_tokens} != {golden.prompt_tokens}" + ) + assert row.completion_tokens == golden.completion_tokens, ( + f"completion_tokens {row.completion_tokens} != {golden.completion_tokens}" ) cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py deleted file mode 100644 index a36bb1a8662..00000000000 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ /dev/null @@ -1,368 +0,0 @@ -"""Wire-format e2e: one scripted upstream per provider wire, answering with a -usage payload where every token kind the wire can report is nonzero. The spend -row's gross input cost must equal fresh tokens at the input rate plus each cache -and audio component at its own rate -- proving the wire's usage shape landed the -cached tokens inside the total (OpenAI/Gemini) or as separate fields -(Anthropic), and that the biller subtracted them before billing fresh tokens. - -Also covers the Responses API wire (an openai/gpt-5.5-pro deployment bridged by -the proxy to POST /responses) and a streamed Anthropic-messages case. -""" - -from __future__ import annotations - -import pytest -from collections.abc import Mapping -from types import MappingProxyType -from typing import Final - -from conftest import CostCalcClient, cost_rows, register_scenario_deployment -from cost_matrix import ( - FRONTIER_MODELS, - Case, - FrontierModel, - expected_breakdown, - expected_token_columns, -) -from e2e_config import unique_marker -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions, ChatTool, ChatToolFunction -from scripted_provider import ScriptedUsage - -pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark - -_MODELS: Final[Mapping[str, FrontierModel]] = MappingProxyType( - {model.map_key: model for model in FRONTIER_MODELS} -) - -# One scripted usage per wire, every reportable token kind nonzero. -_WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ - "openai_chat": ( - "gpt-5.6", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "openai_responses": ( - "gpt-5.5-pro", - ScriptedUsage( - fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25, reasoning_tokens=15 - ), - ), - "anthropic_messages": ( - "claude-sonnet-5", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - ), - ), - "gemini_generate": ( - "gemini/gemini-3.8-flash", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "together_chat": ( - "together_ai/moonshotai/Kimi-K3", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "fireworks_chat": ( - "fireworks_ai/kimi-k3", - ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), - ), - "azure_chat": ( - "azure/gpt-5.6", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "bedrock_converse": ( - "anthropic.claude-sonnet-5-v1:0", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - ), - ), - "vertex_generate": ( - "gemini-3.8-flash", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), -}) - -_SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) - -# Renderer-level shapes the pricing matrix gates per cap, pinned here once per -# wire so the sidecar emits prove they survive the proxy end to end. -_SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( - *( - ( - f"tool_call_{'stream' if stream else 'sync'}", - wire, - Case(name="tool_call", usage=_SHAPE_USAGE, stream=stream, tool_call=True), - ) - for wire in _WIRE_USAGE - for stream in (False, True) - ), - ( - "responses_incomplete", - "openai_responses", - Case(name="stream_no_usage_incomplete", usage=_SHAPE_USAGE, stream=True, terminal="incomplete"), - ), - ( - "responses_unvalidated", - "openai_responses", - Case(name="stream_unvalidated", usage=_SHAPE_USAGE, stream=True, terminal="unvalidated"), - ), - ( - "gemini_prompt_blocked", - "gemini_generate", - Case( - name="prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "gemini_prompt_blocked_stream", - "gemini_generate", - Case( - name="stream_prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "vertex_prompt_blocked", - "vertex_generate", - Case( - name="prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "vertex_prompt_blocked_stream", - "vertex_generate", - Case( - name="stream_prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "azure_served_model_override", - "azure_chat", - Case( - name="response_model_override", - usage=_SHAPE_USAGE, - response_model_override=True, - ), - ), -) - - -def _shape_id(entry: tuple[str, str, Case]) -> str: - return entry[0] - - -class TestWireFormats: - @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_wire_usage_shape_bills_each_component( - self, - client: CostCalcClient, - resources: ResourceManager, - scoped_key: str, - wire: str, - ) -> None: - map_key, usage = _WIRE_USAGE[wire] - model: Final = _MODELS[map_key] - case: Final = Case(name="basic", usage=usage) - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted wire call"),), - ), - ) - assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, f"{wire}: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"{wire}: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - breakdown: Final = row.breakdown - assert breakdown.input_cost is not None and cost_rows.approx_equal( - breakdown.input_cost, expected.input_cost - ), ( - f"{wire}: gross input_cost {breakdown.input_cost} != expected {expected.input_cost}; " - "cached/written tokens billed at the input rate" - ) - assert breakdown.output_cost is not None and cost_rows.approx_equal( - breakdown.output_cost, expected.output_cost - ), f"{wire}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" - - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"{wire}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" - ) - assert row.completion_tokens == completion_tokens, ( - f"{wire}: completion_tokens {row.completion_tokens} != {completion_tokens}" - ) - cost_rows.assert_total_is_sum_of_components(row) - - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_anthropic_streamed_usage_bills_each_component( - self, client: CostCalcClient, resources: ResourceManager, scoped_key: str - ) -> None: - map_key, usage = _WIRE_USAGE["anthropic_messages"] - model: Final = _MODELS[map_key] - case: Final = Case(name="stream", usage=usage, stream=True) - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted anthropic stream"),), - stream=True, - stream_options=ChatStreamOptions(include_usage=True), - ), - stream=True, - ) - assert response.ok, f"anthropic stream: proxy returned {response.status_code}: {response.body[:400]}" - assert response.stream_done, "anthropic stream did not reach its terminal event" - assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, "anthropic stream: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"anthropic stream: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - cost_rows.assert_total_is_sum_of_components(row) - - @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=_shape_id) - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_response_shape_bills_reported_usage( - self, - client: CostCalcClient, - resources: ResourceManager, - scoped_key: str, - shape_wire_case: tuple[str, str, Case], - ) -> None: - shape, wire, case = shape_wire_case - map_key, _usage = _WIRE_USAGE[wire] - model: Final = _MODELS[map_key] - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted {shape}"),), - stream=case.stream, - stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, - tools=( - ( - ChatTool( - function=ChatToolFunction( - name="get_weather", - parameters={"type": "object", "properties": {"city": {"type": "string"}}}, - ) - ), - ) - if case.tool_call - else None - ), - ), - stream=case.stream, - ) - assert response.ok, f"{shape}: proxy returned {response.status_code}: {response.body[:400]}" - if case.stream: - assert response.stream_done, f"{shape}: stream did not reach its terminal event" - assert response.stream_error is None, f"{shape}: stream error: {response.stream_error}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, f"{shape}: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"{shape}: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"{shape}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" - ) - assert row.completion_tokens == completion_tokens, ( - f"{shape}: completion_tokens {row.completion_tokens} != {completion_tokens}" - ) - cost_rows.assert_total_is_sum_of_components(row) From 522a7f569283b9a3bc0fed2a66f1c7545f30b12b Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:51:43 +0000 Subject: [PATCH 043/349] test(e2e): gate all_components cases by rates and tidy cost matrix names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 26 ++++++++ tests/e2e/cost_calculation/cost_matrix.py | 31 +++++---- tests/e2e/cost_calculation/expected.json | 7 --- .../e2e/cost_calculation/generate_expected.py | 63 ++++++++++--------- .../e2e/cost_calculation/test_matrix_data.py | 11 ++-- 5 files changed, 79 insertions(+), 59 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index e898557ea35..3dc4fc4d99c 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -180,11 +180,20 @@ "audio_input_tokens": 5, "audio_output_tokens": 3 }, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "output_cost_per_reasoning_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token" + ], "wires": ["openai_chat", "azure_chat", "together_chat"] }, { "name": "all_components_fireworks", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, + "requires_rates": ["cache_read_input_token_cost"], "wires": ["fireworks_chat"] }, { @@ -196,6 +205,11 @@ "cache_write_1h_tokens": 10, "output_tokens": 25 }, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr" + ], "wires": ["anthropic_messages", "bedrock_converse"] }, { @@ -208,6 +222,11 @@ "output_tokens": 25 }, "stream": true, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr" + ], "wires": ["anthropic_messages"] }, { @@ -220,11 +239,18 @@ "audio_input_tokens": 5, "audio_output_tokens": 3 }, + "requires_rates": [ + "cache_read_input_token_cost", + "output_cost_per_reasoning_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token" + ], "wires": ["gemini_generate", "vertex_generate"] }, { "name": "all_components_responses", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, + "requires_rates": ["cache_read_input_token_cost", "output_cost_per_reasoning_token"], "wires": ["openai_responses"] } ] diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index b03c851d208..a8f60b79ae7 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -27,7 +27,6 @@ from types import MappingProxyType from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter - from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" @@ -69,9 +68,9 @@ class CostMapEntry(BaseModel): web_search_billing_unit: str | None = None -_COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) -_COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( - _COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) +COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( + COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) ) TIER_THRESHOLD_TOKENS: Final = 200_000 @@ -146,10 +145,10 @@ class _CasesFile(BaseModel): cases: tuple[Case, ...] = () -_CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) -CASES: Final[tuple[Case, ...]] = _CASES_FILE.cases +CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) +CASES: Final[tuple[Case, ...]] = CASES_FILE.cases _DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( - {spec.map_key: spec for spec in _CASES_FILE.deployments} + {spec.map_key: spec for spec in CASES_FILE.deployments} ) @@ -221,13 +220,13 @@ class FrontierModel: @property def rates(self) -> CostMapEntry: - return _COST_MAP[self.map_key] + return COST_MAP[self.map_key] @property def override_rates(self) -> CostMapEntry: if self.base_model is not None or self.override_map_key is None: return self.rates - return _COST_MAP[self.override_map_key] + return COST_MAP[self.override_map_key] @property def provider_model(self) -> str: @@ -262,13 +261,13 @@ def _litellm_model_for(map_key: str, wiring: _ProviderWiring) -> str: def _frontier() -> tuple[FrontierModel, ...]: groups: Final[Mapping[tuple[str, str], tuple[str, ...]]] = MappingProxyType( { - pair: tuple(sorted(k for k, e in _COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) - for pair in {(e.litellm_provider, e.mode) for e in _COST_MAP.values()} + pair: tuple(sorted(k for k, e in COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) + for pair in {(e.litellm_provider, e.mode) for e in COST_MAP.values()} } ) models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple - for map_key in sorted(_COST_MAP): - entry: Final = _COST_MAP[map_key] + for map_key in sorted(COST_MAP): + entry: Final = COST_MAP[map_key] pair: Final = (entry.litellm_provider, entry.mode) wiring: Final = _PROVIDER_WIRING.get(pair) if wiring is None: @@ -416,7 +415,7 @@ def image_input_data_url() -> str: IMAGE_INPUT_DATA_URL: Final = image_input_data_url() -class _ExpectedCell(BaseModel): +class ExpectedCell(BaseModel): model_config = ConfigDict(frozen=True) spend: float @@ -426,8 +425,8 @@ class _ExpectedCell(BaseModel): completion_tokens: int -_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, _ExpectedCell]) -EXPECTED: Final[Mapping[str, _ExpectedCell]] = MappingProxyType( +_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, ExpectedCell]) +EXPECTED: Final[Mapping[str, ExpectedCell]] = MappingProxyType( _EXPECTED_ADAPTER.validate_python(json.loads(EXPECTED_PATH.read_text())) if EXPECTED_PATH.exists() else {} diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index 7a92fb2476f..3b18e9ed9f4 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -1686,13 +1686,6 @@ "prompt_tokens": 100, "spend": 0.0216 }, - "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { - "completion_tokens": 25, - "input_cost": 0.020900000000000002, - "output_cost": 0.0095, - "prompt_tokens": 150, - "spend": 0.030400000000000003 - }, "meta.llama4-maverick-17b-instruct-v1:0|basic": { "completion_tokens": 40, "input_cost": 0.0228, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index de979f272fe..c093ecbe0ea 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -14,8 +14,10 @@ from __future__ import annotations import json import sys +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -27,6 +29,7 @@ from cost_matrix import ( # noqa: E402 # path bootstrap before package-local i TIER_THRESHOLD_TOKENS, Case, CostMapEntry, + ExpectedCell, FrontierModel, cases_for, expected_key, @@ -93,16 +96,11 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: or rates.output_cost_per_token or 0.0 ) - # The biller charges cache writes at the input rate when the entry carries - # no cache_creation rate (cost_calculator.py:2452), and at the 5m write - # rate when the 1h variant is unset; cache reads bill only at their own - # rate (zero when the entry lacks one). - write_5m_rate: Final = rates.cache_creation_input_token_cost or in_rate input_cost: Final = ( u.fresh_input_tokens * in_rate + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * write_5m_rate - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or write_5m_rate) + + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) output_cost: Final = ( @@ -147,39 +145,46 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: ) -def _proposed() -> dict[str, dict[str, object]]: - return { - expected_key(model, case): ( - lambda breakdown, tokens: { - "spend": breakdown.total, - "input_cost": breakdown.input_cost, - "output_cost": breakdown.output_cost, - "prompt_tokens": tokens[0], - "completion_tokens": tokens[1], - } - )(expected_breakdown(model, case), expected_token_columns(model, case)) - for model in FRONTIER_MODELS - for case in cases_for(model) - if case.exact_spend - } +def _cell(model: FrontierModel, case: Case) -> ExpectedCell: + breakdown: Final = expected_breakdown(model, case) + prompt_tokens, completion_tokens = expected_token_columns(model, case) + return ExpectedCell( + spend=breakdown.total, + input_cost=breakdown.input_cost, + output_cost=breakdown.output_cost, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + + +def _proposed() -> Mapping[str, ExpectedCell]: + return MappingProxyType( + { + expected_key(model, case): _cell(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + ) def main() -> None: rewrite: Final = "--rewrite" in sys.argv[1:] proposed: Final = _proposed() + proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()} existing: Final = ( json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} ) merged: Final = { - key: (proposed[key] if rewrite or key not in existing else existing[key]) - for key in sorted(proposed) + key: (proposed_values[key] if rewrite or key not in existing else existing[key]) + for key in sorted(proposed_values) } - added: Final = sum(1 for key in proposed if key not in existing) - removed: Final = sum(1 for key in existing if key not in proposed) - kept: Final = sum(1 for key in proposed if key in existing and not rewrite) - rewritten: Final = sum(1 for key in proposed if key in existing and rewrite) + added: Final = sum(1 for key in proposed_values if key not in existing) + removed: Final = sum(1 for key in existing if key not in proposed_values) + kept: Final = sum(1 for key in proposed_values if key in existing and not rewrite) + rewritten: Final = sum(1 for key in proposed_values if key in existing and rewrite) EXPECTED_PATH.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") - print( + print( # noqa: T201 # CLI summary is the tool output f"expected.json: {added} added, {removed} removed, {kept} kept, " f"{rewritten} rewritten ({len(merged)} cells)" ) diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py index fdbb6ddd293..8340257939c 100644 --- a/tests/e2e/cost_calculation/test_matrix_data.py +++ b/tests/e2e/cost_calculation/test_matrix_data.py @@ -8,11 +8,10 @@ from __future__ import annotations from typing import Final import pytest - from cost_matrix import ( - _CASES_FILE, - _COST_MAP, CASES, + CASES_FILE, + COST_MAP, EXPECTED, FRONTIER_MODELS, CostMapEntry, @@ -41,7 +40,7 @@ def test_expected_keys_match_derived_exact_cells() -> None: def test_deployments_reference_existing_map_keys() -> None: unknown: Final = sorted( - spec.map_key for spec in _CASES_FILE.deployments if spec.map_key not in _COST_MAP + spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP ) assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" @@ -55,9 +54,7 @@ def test_requires_rates_are_cost_map_fields() -> None: def test_no_two_entries_share_input_rate() -> None: - rates: Final = [ - entry.input_cost_per_token for entry in _COST_MAP.values() - ] + rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) assert len(rates) == len(set(rates)), ( "two cost_map entries share input_cost_per_token; the suite relies on " "distinct rates so a wrong-model bill can never coincidentally match" From e1c9ae5ae45e3b66041a25a6ae6ca9c7633944b7 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:56:39 +0000 Subject: [PATCH 044/349] test(e2e): drop needless sys.path bootstrap from golden generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/generate_expected.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index c093ecbe0ea..e243e477839 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -16,14 +16,10 @@ import json import sys from collections.abc import Mapping from dataclasses import dataclass -from pathlib import Path from types import MappingProxyType from typing import Final -sys.path.insert(0, str(Path(__file__).resolve().parent)) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) - -from cost_matrix import ( # noqa: E402 # path bootstrap before package-local imports +from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, TIER_THRESHOLD_TOKENS, From 3e11c986766ed7a32ead704e5284fbfeaf889c6b Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:02:41 +0000 Subject: [PATCH 045/349] test(e2e): satisfy pyright in cost matrix derivation and golden generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 14 +++++++------- tests/e2e/cost_calculation/generate_expected.py | 16 +++++++++++++--- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index a8f60b79ae7..35344a099be 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -267,23 +267,23 @@ def _frontier() -> tuple[FrontierModel, ...]: ) models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple for map_key in sorted(COST_MAP): - entry: Final = COST_MAP[map_key] - pair: Final = (entry.litellm_provider, entry.mode) - wiring: Final = _PROVIDER_WIRING.get(pair) + entry = COST_MAP[map_key] + pair = (entry.litellm_provider, entry.mode) + wiring = _PROVIDER_WIRING.get(pair) if wiring is None: raise ValueError( f"cost_map entry {map_key} has no wiring for " f"(litellm_provider={pair[0]}, mode={pair[1]}); add a " f"_ProviderWiring row in cost_matrix.py" ) - siblings: Final = groups[pair] - override_key: Final = ( + siblings = groups[pair] + override_key = ( siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None ) - override_litellm: Final = ( + override_litellm = ( _litellm_model_for(override_key, wiring) if override_key is not None else None ) - deployment: Final = _DEPLOYMENTS.get(map_key) + deployment = _DEPLOYMENTS.get(map_key) models.append( FrontierModel( model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index e243e477839..f5514092c88 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -19,6 +19,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, @@ -168,11 +170,19 @@ def main() -> None: rewrite: Final = "--rewrite" in sys.argv[1:] proposed: Final = _proposed() proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()} - existing: Final = ( - json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} + existing: Final[Mapping[str, ExpectedCell]] = ( + TypeAdapter(dict[str, ExpectedCell]).validate_python( + json.loads(EXPECTED_PATH.read_text()) + ) + if EXPECTED_PATH.exists() + else {} ) merged: Final = { - key: (proposed_values[key] if rewrite or key not in existing else existing[key]) + key: ( + proposed_values[key] + if rewrite or key not in existing + else existing[key].model_dump() + ) for key in sorted(proposed_values) } added: Final = sum(1 for key in proposed_values if key not in existing) From fc0cce553a631e912e7892893bde188e9716b415 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:13:36 +0000 Subject: [PATCH 046/349] test(e2e): derive cache rates from first principles and ungate all_components cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 17 +---------------- tests/e2e/cost_calculation/expected.json | 7 +++++++ tests/e2e/cost_calculation/generate_expected.py | 17 ++++++++++++++--- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index 3dc4fc4d99c..e01ac97e9ff 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -181,9 +181,6 @@ "audio_output_tokens": 3 }, "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", "output_cost_per_reasoning_token", "input_cost_per_audio_token", "output_cost_per_audio_token" @@ -193,7 +190,6 @@ { "name": "all_components_fireworks", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, - "requires_rates": ["cache_read_input_token_cost"], "wires": ["fireworks_chat"] }, { @@ -205,11 +201,6 @@ "cache_write_1h_tokens": 10, "output_tokens": 25 }, - "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr" - ], "wires": ["anthropic_messages", "bedrock_converse"] }, { @@ -222,11 +213,6 @@ "output_tokens": 25 }, "stream": true, - "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr" - ], "wires": ["anthropic_messages"] }, { @@ -240,7 +226,6 @@ "audio_output_tokens": 3 }, "requires_rates": [ - "cache_read_input_token_cost", "output_cost_per_reasoning_token", "input_cost_per_audio_token", "output_cost_per_audio_token" @@ -250,7 +235,7 @@ { "name": "all_components_responses", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, - "requires_rates": ["cache_read_input_token_cost", "output_cost_per_reasoning_token"], + "requires_rates": ["output_cost_per_reasoning_token"], "wires": ["openai_responses"] } ] diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index 3b18e9ed9f4..caea2c3c764 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -1686,6 +1686,13 @@ "prompt_tokens": 100, "spend": 0.0216 }, + "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.0285, + "output_cost": 0.0095, + "prompt_tokens": 150, + "spend": 0.038 + }, "meta.llama4-maverick-17b-instruct-v1:0|basic": { "completion_tokens": 40, "input_cost": 0.0228, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index f5514092c88..a6eabcc7286 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -94,11 +94,22 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: or rates.output_cost_per_token or 0.0 ) + write_rate: Final = ( + rates.cache_creation_input_token_cost + if rates.cache_creation_input_token_cost is not None + else in_rate + ) input_cost: Final = ( u.fresh_input_tokens * in_rate - + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + + u.cache_read_tokens + * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate) + + u.cache_write_5m_tokens * write_rate + + u.cache_write_1h_tokens + * ( + rates.cache_creation_input_token_cost_above_1hr + if rates.cache_creation_input_token_cost_above_1hr is not None + else write_rate + ) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) output_cost: Final = ( From 072b32baf2097e5421672956ce34809106f592aa Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:18:02 +0000 Subject: [PATCH 047/349] test(e2e): derive goldens from first-principles rate selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 10 ++- tests/e2e/cost_calculation/expected.json | 18 ++-- .../e2e/cost_calculation/generate_expected.py | 86 +++++++++---------- 3 files changed, 61 insertions(+), 53 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index e01ac97e9ff..49eebc85231 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -62,7 +62,15 @@ "name": "web_search", "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3}, "requires_rates": ["search_context_cost_per_query"], - "requires_caps": ["web_search"] + "requires_caps": ["web_search"], + "wires": ["openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"] + }, + { + "name": "web_search_single", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 1}, + "requires_rates": ["search_context_cost_per_query"], + "requires_caps": ["web_search"], + "wires": ["openai_chat", "together_chat", "fireworks_chat", "azure_chat"] }, { "name": "stream", diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index caea2c3c764..984b670a82c 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -160,7 +160,7 @@ "prompt_tokens": 120, "spend": 0.032 }, - "azure/gpt-5.4-mini|web_search": { + "azure/gpt-5.4-mini|web_search_single": { "completion_tokens": 30, "input_cost": 0.016, "output_cost": 0.009600000000000001, @@ -272,7 +272,7 @@ "prompt_tokens": 120, "spend": 0.03 }, - "azure/gpt-5.6|web_search": { + "azure/gpt-5.6|web_search_single": { "completion_tokens": 30, "input_cost": 0.015, "output_cost": 0.009, @@ -615,7 +615,7 @@ "prompt_tokens": 120, "spend": 0.028000000000000004 }, - "fireworks_ai/deepseek-v4p1-flash|web_search": { + "fireworks_ai/deepseek-v4p1-flash|web_search_single": { "completion_tokens": 30, "input_cost": 0.014000000000000002, "output_cost": 0.008400000000000001, @@ -706,7 +706,7 @@ "prompt_tokens": 120, "spend": 0.024 }, - "fireworks_ai/kimi-k3|web_search": { + "fireworks_ai/kimi-k3|web_search_single": { "completion_tokens": 30, "input_cost": 0.012000000000000002, "output_cost": 0.007200000000000001, @@ -797,7 +797,7 @@ "prompt_tokens": 120, "spend": 0.026000000000000002 }, - "fireworks_ai/qwen3p8-max|web_search": { + "fireworks_ai/qwen3p8-max|web_search_single": { "completion_tokens": 30, "input_cost": 0.013000000000000001, "output_cost": 0.007800000000000001, @@ -1462,7 +1462,7 @@ "prompt_tokens": 120, "spend": 0.008 }, - "gpt-5.4-mini|web_search": { + "gpt-5.4-mini|web_search_single": { "completion_tokens": 30, "input_cost": 0.004, "output_cost": 0.0024000000000000002, @@ -1679,7 +1679,7 @@ "prompt_tokens": 120, "spend": 0.002 }, - "gpt-5.6|web_search": { + "gpt-5.6|web_search_single": { "completion_tokens": 30, "input_cost": 0.001, "output_cost": 0.0006000000000000001, @@ -1826,7 +1826,7 @@ "prompt_tokens": 120, "spend": 0.02 }, - "together_ai/moonshotai/Kimi-K3|web_search": { + "together_ai/moonshotai/Kimi-K3|web_search_single": { "completion_tokens": 30, "input_cost": 0.01, "output_cost": 0.006, @@ -1938,7 +1938,7 @@ "prompt_tokens": 120, "spend": 0.022 }, - "together_ai/zai-org/GLM-5.3|web_search": { + "together_ai/zai-org/GLM-5.3|web_search_single": { "completion_tokens": 30, "input_cost": 0.011000000000000001, "output_cost": 0.0066, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index a6eabcc7286..64abdb14c99 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -19,8 +19,6 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final -from pydantic import TypeAdapter - from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, @@ -32,19 +30,11 @@ from cost_matrix import ( cases_for, expected_key, ) - -# Wires whose response surface reports a real web-search call count; the -# chat-completions wires only expose url_citation annotations, so their billed -# count floors to one. -_EXACT_WEB_SEARCH_WIRES: Final = frozenset( - {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"} -) +from pydantic import TypeAdapter -def billed_web_search_calls(model: FrontierModel, case: Case) -> int: - if case.usage.web_search_calls == 0: - return 0 - return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1 +def _first_present(*rates: float | None) -> float | None: + return next((rate for rate in rates if rate is not None), None) @dataclass(frozen=True, slots=True) @@ -67,11 +57,14 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; output = text*out + reasoning*reasoning + audio_out*audio_out; plus the - billed web-search calls at the medium search-context rate. Above-threshold - swaps every input/output rate to its ``_above_200k_tokens`` variant when - total prompt tokens exceed the threshold; a service tier swaps input/output - to the tier's variants, falling back to the base rate when a variant is - unset -- mirroring _get_token_base_cost in litellm's cost calculator. + billed web-search calls at the medium search-context rate. Every billed + token is a token the provider charged for: a component whose entry has no + dedicated rate bills at the ordinary input or output rate, and a present + rate (including an explicit 0.0) is authoritative. When the total prompt + tokens exceed the threshold, input/output rates come from the + ``_above_200k_tokens`` variants; a service tier takes its ``_priority`` or + ``_flex`` variant when the entry carries one, and otherwise bills at the + base rate. """ rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates u: Final = case.usage @@ -81,46 +74,53 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: ) tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS in_rate: Final = ( - (rates.input_cost_per_token_above_200k_tokens if tiered else None) - or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.input_cost_per_token + _first_present( + rates.input_cost_per_token_above_200k_tokens if tiered else None, + rates.input_cost_per_token_priority if case.service_tier == "priority" else None, + rates.input_cost_per_token_flex if case.service_tier == "flex" else None, + rates.input_cost_per_token, + ) or 0.0 ) out_rate: Final = ( - (rates.output_cost_per_token_above_200k_tokens if tiered else None) - or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.output_cost_per_token + _first_present( + rates.output_cost_per_token_above_200k_tokens if tiered else None, + rates.output_cost_per_token_priority if case.service_tier == "priority" else None, + rates.output_cost_per_token_flex if case.service_tier == "flex" else None, + rates.output_cost_per_token, + ) or 0.0 ) - write_rate: Final = ( - rates.cache_creation_input_token_cost - if rates.cache_creation_input_token_cost is not None - else in_rate + read_rate: Final = _first_present(rates.cache_read_input_token_cost, in_rate) or 0.0 + write_rate: Final = _first_present(rates.cache_creation_input_token_cost, in_rate) or 0.0 + write_1h_rate: Final = ( + _first_present(rates.cache_creation_input_token_cost_above_1hr, write_rate) or 0.0 ) + audio_in_rate: Final = _first_present(rates.input_cost_per_audio_token, in_rate) or 0.0 + reasoning_rate: Final = _first_present(rates.output_cost_per_reasoning_token, out_rate) or 0.0 + audio_out_rate: Final = _first_present(rates.output_cost_per_audio_token, out_rate) or 0.0 input_cost: Final = ( u.fresh_input_tokens * in_rate - + u.cache_read_tokens - * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate) + + u.cache_read_tokens * read_rate + u.cache_write_5m_tokens * write_rate - + u.cache_write_1h_tokens - * ( - rates.cache_creation_input_token_cost_above_1hr - if rates.cache_creation_input_token_cost_above_1hr is not None - else write_rate - ) - + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + + u.cache_write_1h_tokens * write_1h_rate + + u.audio_input_tokens * audio_in_rate ) output_cost: Final = ( u.output_tokens * out_rate - + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) - + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + + u.reasoning_tokens * reasoning_rate + + u.audio_output_tokens * audio_out_rate ) search: Final = rates.search_context_cost_per_query - tool_cost: Final = billed_web_search_calls(model, case) * ( - search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + medium_rate: Final = ( + search.search_context_size_medium if search is not None else None ) + if u.web_search_calls and medium_rate is None: + raise ValueError( + f"{model.map_key}: case {case.name} bills {u.web_search_calls} web-search " + "calls but the entry has no search_context_cost_per_query medium rate" + ) + tool_cost: Final = u.web_search_calls * (medium_rate if medium_rate is not None else 0.0) return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) From 1de633ac36644c5774cf629a793b6716a98b7580 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:22:01 +0000 Subject: [PATCH 048/349] test(e2e): move matrix data freshness checks to collection time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/cost_calculation/cost_matrix.py | 48 +++++++++++++++ .../e2e/cost_calculation/test_matrix_data.py | 61 ------------------- .../test_token_pricing_e2e.py | 4 ++ 4 files changed, 53 insertions(+), 62 deletions(-) delete mode 100644 tests/e2e/cost_calculation/test_matrix_data.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 707d35b4aa6..f89b3203622 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`; `cost_matrix.matrix_data_errors()` runs at collection time so a stale key set fails the suite's collection loudly), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 35344a099be..68f3186809d 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -435,3 +435,51 @@ EXPECTED: Final[Mapping[str, ExpectedCell]] = MappingProxyType( def expected_key(model: FrontierModel, case: Case) -> str: return f"{model.map_key}|{case.name}" + + +def matrix_data_errors() -> tuple[str, ...]: + """Freshness findings for the data files, as human-readable strings. + + Called at collection time by the e2e suite; also usable from + generate_expected.py's context without importing pytest. + """ + derived: Final = { + expected_key(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + golden: Final = set(EXPECTED) + unknown_deployments: Final = sorted( + spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP + ) + unknown_rates: Final = sorted( + {field for case in CASES for field in case.requires_rates} - set(CostMapEntry.model_fields) + ) + input_rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) + findings: Final = ( + ( + "expected.json is out of sync with the derived matrix; run " + "uv run python tests/e2e/cost_calculation/generate_expected.py " + f"(missing: {sorted(derived - golden)}; stale: {sorted(golden - derived)})" + ) + if derived != golden + else None, + ( + f"deployments entries name map keys absent from cost_map.json: {unknown_deployments}" + if unknown_deployments + else None + ), + ( + f"requires_rates names that are not CostMapEntry fields: {unknown_rates}" + if unknown_rates + else None + ), + ( + "two cost_map entries share input_cost_per_token; the suite relies on " + "distinct rates so a wrong-model bill can never coincidentally match" + if len(input_rates) != len(set(input_rates)) + else None + ), + ) + return tuple(finding for finding in findings if finding is not None) diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py deleted file mode 100644 index 8340257939c..00000000000 --- a/tests/e2e/cost_calculation/test_matrix_data.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Freshness checks for the cost suite's data files; markerless, so it runs on -any pytest invocation of the folder without the stack. expected.json is the -oracle: these tests check its key set against the derived matrix, never its -values (the generator proposes, the file decides).""" - -from __future__ import annotations - -from typing import Final - -import pytest -from cost_matrix import ( - CASES, - CASES_FILE, - COST_MAP, - EXPECTED, - FRONTIER_MODELS, - CostMapEntry, - cases_for, - expected_key, -) - - -def test_expected_keys_match_derived_exact_cells() -> None: - derived: Final = { - expected_key(model, case) - for model in FRONTIER_MODELS - for case in cases_for(model) - if case.exact_spend - } - golden: Final = set(EXPECTED) - if derived != golden: - missing: Final = sorted(derived - golden) - stale: Final = sorted(golden - derived) - pytest.fail( - "expected.json is out of sync with the derived matrix; run " - "uv run python tests/e2e/cost_calculation/generate_expected.py " - f"(missing: {missing}; stale: {stale})" - ) - - -def test_deployments_reference_existing_map_keys() -> None: - unknown: Final = sorted( - spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP - ) - assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" - - -def test_requires_rates_are_cost_map_fields() -> None: - fields: Final = set(CostMapEntry.model_fields) - unknown: Final = sorted( - {field for case in CASES for field in case.requires_rates} - fields - ) - assert not unknown, f"requires_rates names that are not CostMapEntry fields: {unknown}" - - -def test_no_two_entries_share_input_rate() -> None: - rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) - assert len(rates) == len(set(rates)), ( - "two cost_map entries share input_cost_per_token; the suite relies on " - "distinct rates so a wrong-model bill can never coincidentally match" - ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 7cd128ad6fb..346a55aa22d 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -22,6 +22,7 @@ from cost_matrix import ( FrontierModel, cases_for, expected_key, + matrix_data_errors, recount_cost, ) from e2e_config import unique_marker @@ -39,6 +40,9 @@ from models import ( pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark +if _data_errors := matrix_data_errors(): + raise ValueError("\n".join(_data_errors)) + _MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple( (model, case) for model in FRONTIER_MODELS for case in cases_for(model) ) From ef1f306a7dc0777276166859b3fa4d2e6272cefb Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:53:52 +0000 Subject: [PATCH 049/349] test(e2e): emit gemini stream usage only on the final chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/scripted_provider.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 982132ed8df..90d95441e5c 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -750,10 +750,8 @@ def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, objec def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: emit_usage: Final = scenario.stream_usage == "final_chunk" - first: Final = ( - _jobj(*((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata")) - if scenario.stream_usage == "absent" - else _gemini_body(scenario, requested_model) + first: Final = _jobj( + *((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata") ) return _sse( ( From e254377049ca6f7087693cff4b99b291c9ba6010 Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:07:57 +0100 Subject: [PATCH 050/349] fix(azure): drop tool_choice without tools DEVX-829 --- litellm/llms/azure/chat/gpt_transformation.py | 9 +- .../test_azure_chat_gpt_transformation.py | 123 ++++++++++++++++++ ...test_azure_chat_o_series_transformation.py | 3 +- 3 files changed, 133 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 6d17a1359bc..0debbe4f74d 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,10 +280,17 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) + request_params: Final = { + key: value + for key, value in optional_params.items() + if key != "tool_choice" + or optional_params.get("tools") + or optional_params.get("functions") + } return { "model": model, "messages": azure_messages, - **optional_params, + **request_params, **sanitized_tools_update(optional_params), } diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index e8b98c696e1..92a8124d8a9 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -333,3 +333,126 @@ class TestAzureToolSchemaCombinatorFlattening: ) assert "tools" not in request assert request["temperature"] == 0.2 + + +@pytest.mark.parametrize("tool_choice", ["none", "auto"]) +def test_azure_drops_tool_choice_without_tools_or_functions(tool_choice: str) -> None: + optional_params = {"tool_choice": tool_choice, "temperature": 0.2} + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert "tool_choice" not in request + assert request["temperature"] == 0.2 + assert optional_params["tool_choice"] == tool_choice + + +def test_azure_tools_empty_drops_tool_choice() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": [], "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["tools"] == [] + assert "tool_choice" not in request + + +def test_azure_functions_empty_drops_tool_choice() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"functions": [], "tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["functions"] == [] + assert "tool_choice" not in request + + +def test_azure_preserves_tool_choice_with_tools() -> None: + tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["tools"] == tools + assert request["tool_choice"] == "auto" + + +def test_azure_preserves_tool_choice_with_legacy_functions() -> None: + functions = [{"name": "get_weather", "parameters": {}}] + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"functions": functions, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["functions"] == functions + assert request["tool_choice"] == "auto" + + +def test_azure_preserves_function_call_without_tools() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"function_call": "none", "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["function_call"] == "none" + assert "tool_choice" not in request + + +def test_azure_gpt5_drops_tool_choice_without_tools() -> None: + request = AzureOpenAIGPT5Config().transform_request( + model="gpt5_series/gpt-5.6-sol", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["model"] == "gpt-5.6-sol" + assert "tool_choice" not in request + + +@pytest.mark.asyncio +async def test_azure_async_transform_drops_tool_choice_without_tools() -> None: + request = await AzureOpenAIConfig().async_transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert "tool_choice" not in request + + +@pytest.mark.asyncio +async def test_azure_gpt5_async_transform_drops_tool_choice_without_tools() -> None: + request = await AzureOpenAIGPT5Config().async_transform_request( + model="gpt5_series/gpt-5.6-sol", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["model"] == "gpt-5.6-sol" + assert "tool_choice" not in request diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py index 9db9ab971a0..57d60df3a11 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -14,7 +14,7 @@ async def test_azure_chat_o_series_transformation(): provider_config = AzureOpenAIO1Config() model = "o_series/web-interface-o1-mini" messages = [{"role": "user", "content": "Hello, how are you?"}] - optional_params = {} + optional_params = {"tool_choice": "none"} litellm_params = {} headers = {} @@ -23,6 +23,7 @@ async def test_azure_chat_o_series_transformation(): ) print(response) assert response["model"] == "web-interface-o1-mini" + assert "tool_choice" not in response def test_azure_o_series_transform_request_flattens_top_level_anyof(): From 48712f733a641f16b0b4fe60c221fd9f1c7076fa Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:30:06 +0100 Subject: [PATCH 051/349] style(azure): format request parameter filter DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0debbe4f74d..7cb50ee5348 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -283,9 +283,7 @@ class AzureOpenAIConfig(BaseConfig): request_params: Final = { key: value for key, value in optional_params.items() - if key != "tool_choice" - or optional_params.get("tools") - or optional_params.get("functions") + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") } return { "model": model, From bd222bd8d9f6d083b8058c5fef3e998b4f92b3af Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 08:18:36 +0100 Subject: [PATCH 052/349] fix(azure): avoid mutable request mapping DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 7cb50ee5348..424422612db 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,11 +280,13 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) - request_params: Final = { - key: value - for key, value in optional_params.items() - if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") - } + request_params: Final = MappingProxyType( + { + key: value + for key, value in optional_params.items() + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") + } + ) return { "model": model, "messages": azure_messages, From 2fa115db2bac68ac90b0b900dc9eecb728c3bd4d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 17 Sep 2026 12:53:15 -0400 Subject: [PATCH 053/349] feat(router): add maintained Fuse model and harness presets --- .../public_endpoints/public_endpoints.py | 9 + .../complexity_router/README.md | 45 +++++ .../complexity_router/fuse_presets.json | 100 +++++++++++ .../complexity_router/fuse_presets.py | 52 ++++++ .../complexity_router/llm_v2.py | 37 +++- pyproject.toml | 1 + .../public_endpoints/test_public_endpoints.py | 11 ++ .../router_strategy/test_fuse_presets.py | 43 +++++ .../router_strategy/test_llm_v2.py | 110 ++++++++++++ .../test_auto_router_model_naming.py | 49 ++++++ ...ecastClassifierConfig.integration.test.tsx | 163 +++++++++++++++++- .../add_model/ForecastClassifierConfig.tsx | 26 +-- .../add_model/FuseProfilePresets.tsx | 117 +++++++++++++ .../build_complexity_router_config.test.ts | 17 ++ .../forecast_classifier_config.test.ts | 65 ++++++- .../add_model/forecast_classifier_config.ts | 33 +++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 82 ++++++++- 17 files changed, 917 insertions(+), 43 deletions(-) create mode 100644 litellm/router_strategy/complexity_router/fuse_presets.json create mode 100644 litellm/router_strategy/complexity_router/fuse_presets.py create mode 100644 tests/test_litellm/router_strategy/test_fuse_presets.py create mode 100644 ui/litellm-dashboard/src/components/add_model/FuseProfilePresets.tsx diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 94a59828451..e395f56194f 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.utils import get_custom_url from litellm.repositories.table_repositories import ClaudeCodePluginRepository +from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -424,6 +425,14 @@ async def get_complexity_scorer_defaults() -> ComplexityScorerDefaults: ) +@router.get( + "/public/complexity_router/fuse_presets", + response_model=FusePresetCatalog, +) +async def get_public_fuse_presets() -> FusePresetCatalog: + return get_fuse_presets() + + @router.get( "/public/litellm_model_cost_map", tags=["public", "model management"], diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 6505746bca1..d9159cea426 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -179,6 +179,51 @@ Configure capability forecasting through YAML or the model-management API. The dashboard preserves its classifier and calibration on an untouched save; it does not provide a capability-card editor +### Fuse v2 profile presets + +Fuse v2 accepts maintained model and runtime descriptions instead of requiring +custom prose for both solvers and the harness. Select profiles explicitly for +all deployments behind your configured model groups and their actual settings. +Group names do not select profiles automatically + +```yaml +complexity_router_config: + classifier_type: llm_v2 + classifier_llm_config: + model: your-judge-group + tiers: + SIMPLE: your-efficient-group + REASONING: your-capable-group + llm_v2_config: + efficient_profile_preset: claude-sonnet-5-v1 + capable_profile_preset: claude-fable-5-1-v1 + harness_preset: claude-code-v1 + max_quality_gap: 0.05 +``` + +`GET /public/complexity_router/fuse_presets` returns the catalog version, model +profiles, and runtime descriptions, including source URLs. The bundled catalog +is loaded once per process without network requests. Sources are citations only + +Each of `efficient_profile`, `capable_profile`, and `harness` requires either +nonblank custom text or its corresponding preset reference. Custom text wins +when both are supplied, but an unknown or wrong-kind preset is still rejected. +Explicit blank text is invalid even with a valid preset. Custom text remains +limited to 4000 characters + +Saved configurations retain preset references and explicit text separately. +Preset text is resolved when building the classifier prompt, not copied into +stored custom fields. Existing all-custom configurations keep the same prompt. +Versioned preset IDs identify immutable content: revised wording receives a new +ID, and older referenced entries must remain available + +The runtime presets do not imply a repository, runnable tests, network access, +additional tools, or a step, time, or spending budget. mini-SWE-agent describes +an agent interface, not a SWE-bench task. Model descriptions summarize provider +positioning without solve rates or guaranteed rankings. Wording is an evaluation +input, not a calibrated quality claim. Existing Fuse licensing, policy, +calibration, and prompt version are unchanged + ### Heuristic v2 Set `classifier_type: heuristic_v2` to classify with the bundled calibrated diff --git a/litellm/router_strategy/complexity_router/fuse_presets.json b/litellm/router_strategy/complexity_router/fuse_presets.json new file mode 100644 index 00000000000..4006366dc25 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.json @@ -0,0 +1,100 @@ +{ + "version": "2026-09-17-v1", + "models": [ + { + "id": "gpt-6-astra-v1", + "label": "GPT-6 Astra", + "model": "gpt-6-astra", + "text": "OpenAI model for demanding end-to-end work, including reasoning, coding, research, and document tasks", + "sources": ["https://developers.openai.com/api/docs/models/gpt-6-astra"] + }, + { + "id": "gpt-5.6-sol-v1", + "label": "GPT-5.6 Sol", + "model": "gpt-5.6-sol", + "text": "OpenAI model for complex professional work, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-sol"] + }, + { + "id": "gpt-5.6-luna-v1", + "label": "GPT-5.6 Luna", + "model": "gpt-5.6-luna", + "text": "OpenAI model for high-volume workloads, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-luna"] + }, + { + "id": "gpt-5.6-terra-v1", + "label": "GPT-5.6 Terra", + "model": "gpt-5.6-terra", + "text": "OpenAI general-purpose model supporting reasoning, text and image input, and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-terra"] + }, + { + "id": "claude-haiku-4-5-v1", + "label": "Claude Haiku 4.5", + "model": "claude-haiku-4-5", + "text": "Anthropic latency-focused model supporting text and image input, tool use, and extended thinking", + "sources": ["https://platform.claude.com/docs/en/models/haiku-4-5/overview"] + }, + { + "id": "claude-sonnet-5-v1", + "label": "Claude Sonnet 5", + "model": "claude-sonnet-5", + "text": "Anthropic model balancing speed and capability, with adaptive thinking and tool use", + "sources": ["https://platform.claude.com/docs/en/models/sonnet-5/overview"] + }, + { + "id": "claude-opus-5-v1", + "label": "Claude Opus 5", + "model": "claude-opus-5", + "text": "Anthropic model for complex agentic coding and enterprise work, with adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/opus-5/overview"] + }, + { + "id": "claude-fable-5-v1", + "label": "Claude Fable 5", + "model": "claude-fable-5", + "text": "Anthropic model for demanding reasoning and long-running agent tasks, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5/introducing-claude-fable-5-and-claude-mythos-5"] + }, + { + "id": "claude-fable-5-1-v1", + "label": "Claude Fable 5.1", + "model": "claude-fable-5-1", + "text": "Anthropic model for demanding reasoning, long-running agentic coding, and multistep research, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5-1/overview"] + } + ], + "harnesses": [ + { + "id": "unspecified-v1", + "label": "Unspecified runtime", + "text": "Agent runtime is unspecified. Assess the task using the supplied context without assuming repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works", "https://mini-swe-agent.com/latest/faq/"] + }, + { + "id": "claude-code-v1", + "label": "Claude Code", + "text": "Claude Code supplies an agent loop with context management and configured tools. Available actions depend on the session's tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works"] + }, + { + "id": "codex-cli-v1", + "label": "Codex CLI", + "text": "Codex CLI supplies a terminal-based coding agent. File operations, command execution, and integrations depend on the session's tools, permissions, and sandbox. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://learn.chatgpt.com/docs/codex/cli", "https://learn.chatgpt.com/codex/permissions"] + }, + { + "id": "opencode-v1", + "label": "OpenCode", + "text": "OpenCode supplies a configurable agent runtime. Available actions depend on the selected agent, tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://opencode.ai/docs/agents/"] + }, + { + "id": "mini-swe-agent-v1", + "label": "mini-SWE-agent", + "text": "The standard mini-SWE-agent setup uses a bash-only action interface and separate command executions. Available commands and resources depend on its configured environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://mini-swe-agent.com/latest/faq/"] + } + ] +} diff --git a/litellm/router_strategy/complexity_router/fuse_presets.py b/litellm/router_strategy/complexity_router/fuse_presets.py new file mode 100644 index 00000000000..66a96ec5ad5 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.py @@ -0,0 +1,52 @@ +from functools import lru_cache +from importlib.resources import files +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, StringConstraints + +ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] + + +class FuseModelPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + model: str + + +class FuseHarnessPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + + +class FusePresetCatalog(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + version: str + models: tuple[FuseModelPreset, ...] + harnesses: tuple[FuseHarnessPreset, ...] + + +@lru_cache(maxsize=1) +def get_fuse_presets() -> FusePresetCatalog: + return FusePresetCatalog.model_validate_json( + files(__package__).joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + + +def resolve_fuse_profile(text: str | None, preset_id: str | None, kind: Literal["model", "harness"]) -> str | None: + if preset_id is None: + return text + catalog: Final = get_fuse_presets() + presets: Final = catalog.models if kind == "model" else catalog.harnesses + preset: Final = next((entry for entry in presets if entry.id == preset_id), None) + if preset is None: + return None + return text if text is not None else preset.text diff --git a/litellm/router_strategy/complexity_router/llm_v2.py b/litellm/router_strategy/complexity_router/llm_v2.py index 2f545a65aaa..18351237e65 100644 --- a/litellm/router_strategy/complexity_router/llm_v2.py +++ b/litellm/router_strategy/complexity_router/llm_v2.py @@ -7,15 +7,15 @@ from dataclasses import dataclass from sys import float_info from typing import Annotated, Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter, model_validator from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # legacy output validated below ) +from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)] -ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] class _SolverProfile(TypedDict): @@ -139,20 +139,41 @@ class LLMV2Config(BaseModel): efficient_tier: str = "SIMPLE" capable_tier: str = "REASONING" - efficient_profile: ProfileText - capable_profile: ProfileText - harness: ProfileText + efficient_profile: ProfileText | None = None + capable_profile: ProfileText | None = None + harness: ProfileText | None = None + efficient_profile_preset: str | None = None + capable_profile_preset: str | None = None + harness_preset: str | None = None max_quality_gap: float = Field(ge=0.0, le=1.0, description="Maximum estimated success loss allowed for efficient.") max_output_tokens: int = Field(default=1024, ge=1) response_format: Literal["json_schema", "json_object"] = "json_schema" calibration: LLMV2Calibration | None = None + @model_validator(mode="after") + def validate_profiles(self) -> LLMV2Config: + self._profile_texts() + return self + + def _profile_texts(self) -> tuple[str, str, str]: + efficient: Final = resolve_fuse_profile(self.efficient_profile, self.efficient_profile_preset, "model") + capable: Final = resolve_fuse_profile(self.capable_profile, self.capable_profile_preset, "model") + harness: Final = resolve_fuse_profile(self.harness, self.harness_preset, "harness") + if efficient is None: + raise ValueError("efficient_profile requires text or a known efficient_profile_preset") + if capable is None: + raise ValueError("capable_profile requires text or a known capable_profile_preset") + if harness is None: + raise ValueError("harness requires text or a known harness_preset") + return efficient, capable, harness + def system_prompt(self, efficient_model: str, capable_model: str) -> str: + efficient, capable, harness = self._profile_texts() profiles: Final[_SolverProfiles] = { "prompt_version": LLM_V2_PROMPT_VERSION, - "harness": self.harness, - "efficient": {"model": efficient_model, "profile": self.efficient_profile}, - "capable": {"model": capable_model, "profile": self.capable_profile}, + "harness": harness, + "efficient": {"model": efficient_model, "profile": efficient}, + "capable": {"model": capable_model, "profile": capable}, } schema: Final = ( "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) diff --git a/pyproject.toml b/pyproject.toml index 93ff55c4069..2d6d133cd0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -290,6 +290,7 @@ editable-profile = "dev" include = [ "litellm/proxy/_experimental/out/**", "litellm/router_strategy/complexity_router/artifacts/*.json", + "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/client/cli/commands/codex_base_instructions.md", ] exclude = [ diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 0d82ed778f5..b830eb588a5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -11,12 +11,23 @@ from fastapi.testclient import TestClient from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.public_endpoints import router +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) from litellm.types.utils import LlmProviders +def test_fuse_presets_route_serves_the_shared_catalog_without_authentication() -> None: + app: Final = FastAPI() + app.include_router(router) + client: Final = TestClient(app) + response: Final = client.get("/public/complexity_router/fuse_presets") + assert response.status_code == 200 + assert response.json() == get_fuse_presets().model_dump(mode="json") + assert client.get("/public/complexity_router/fuse_presets").json() == response.json() + + def test_get_supported_providers_returns_enum_values(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/test_litellm/router_strategy/test_fuse_presets.py b/tests/test_litellm/router_strategy/test_fuse_presets.py new file mode 100644 index 00000000000..0b8d936383b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_fuse_presets.py @@ -0,0 +1,43 @@ +import json +from importlib.resources import files +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets, resolve_fuse_profile + + +def test_catalog_is_loaded_once_and_preserves_bundled_content() -> None: + get_fuse_presets.cache_clear() + first: Final = get_fuse_presets() + second: Final = get_fuse_presets() + assert first is second + bundled: Final = json.loads( + files("litellm.router_strategy.complexity_router").joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + assert first.model_dump(mode="json") == bundled + entries: Final = (*first.models, *first.harnesses) + assert len({entry.id for entry in entries}) == len(entries) + assert len(first.models) == 9 + assert len(first.harnesses) == 5 + assert all(entry.sources and all(source.startswith("https://") for source in entry.sources) for entry in entries) + + +def test_every_catalog_entry_resolves_without_changing_custom_ownership() -> None: + catalog: Final = get_fuse_presets() + for entry in catalog.models: + assert resolve_fuse_profile(None, entry.id, "model") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "model") == "Custom text" + for entry in catalog.harnesses: + assert resolve_fuse_profile(None, entry.id, "harness") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "harness") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "model") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "harness") == "Custom text" + + +def test_cached_catalog_and_records_cannot_be_modified() -> None: + catalog: Final = get_fuse_presets() + for record, field in ((catalog, "version"), (catalog.models[0], "text"), (catalog.harnesses[0], "text")): + with pytest.raises(ValidationError, match="frozen"): + setattr(record, field, "Changed") diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index 5447c8b43ce..27d31cbe640 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -11,8 +11,10 @@ from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.router_strategy.complexity_router.llm_v2 import ( LLM_V2_PROMPT_VERSION, + LLM_V2_SYSTEM_PROMPT, LLMV2Calibration, LLMV2Config, LLMV2ProbabilityCalibration, @@ -174,6 +176,114 @@ def test_invalid_forecast_settings_are_rejected(overrides: dict[str, object]) -> LLMV2Config.model_validate({**base.model_dump(), **overrides}) +def _preset_config(**overrides: object) -> LLMV2Config: + catalog: Final = get_fuse_presets() + return LLMV2Config.model_validate( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[-1].id, + "max_quality_gap": 0.05, + **overrides, + } + ) + + +def test_preset_roundtrip_keeps_references_without_materializing_text() -> None: + config: Final = _preset_config() + serialized: Final = config.model_dump(exclude_none=True) + assert serialized["efficient_profile_preset"] == config.efficient_profile_preset + assert serialized["capable_profile_preset"] == config.capable_profile_preset + assert serialized["harness_preset"] == config.harness_preset + assert not {"efficient_profile", "capable_profile", "harness"}.intersection(serialized) + assert LLMV2Config.model_validate(config.model_dump()) == config + assert LLMV2Config.model_validate_json(config.model_dump_json()) == config + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_explicit_override_wins_and_survives_roundtrip(field: str) -> None: + config: Final = _preset_config(**{field: " Operator description "}) + roundtrip: Final = LLMV2Config.model_validate_json(config.model_dump_json()) + assert roundtrip.model_dump()[field] == "Operator description" + assert roundtrip.efficient_profile_preset == config.efficient_profile_preset + assert roundtrip.capable_profile_preset == config.capable_profile_preset + assert roundtrip.harness_preset == config.harness_preset + payload: Final = json.loads( + roundtrip.system_prompt("opaque-efficient", "opaque-capable").split("Configured solver profiles:\n")[1] + ) + if field == "harness": + assert payload["harness"] == "Operator description" + else: + assert payload[field.removesuffix("_profile")]["profile"] == "Operator description" + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("invalid", ("", " \n\t", "x" * 4001)) +def test_preset_does_not_bypass_supplied_text_bounds(field: str, invalid: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{field: invalid}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("override", (None, "Custom override")) +@pytest.mark.parametrize("invalid_id", ("missing-v1", "")) +def test_preset_unknown_reference_rejects_even_when_overridden( + field: str, override: str | None, invalid_id: str +) -> None: + with pytest.raises(ValidationError, match=f"{field}.*preset"): + _preset_config(**{field: override, f"{field}_preset": invalid_id}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_missing_text_and_reference_rejects(field: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": None}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_reference_rejects_the_wrong_catalog_kind(field: str) -> None: + catalog: Final = get_fuse_presets() + wrong_id: Final = catalog.models[0].id if field == "harness" else catalog.harnesses[0].id + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": wrong_id}) + + +@pytest.mark.parametrize("mode", ("json_schema", "json_object")) +def test_custom_profile_prompt_bytes_are_unchanged(mode: str) -> None: + base: Final = _config().llm_v2_config + assert base is not None + config: Final = LLMV2Config.model_validate({**base.model_dump(), "response_format": mode}) + old_payload: Final = { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": config.harness, + "efficient": {"model": "opaque-efficient", "profile": config.efficient_profile}, + "capable": {"model": "opaque-capable", "profile": config.capable_profile}, + } + schema: Final = ( + "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) if mode == "json_object" else "" + ) + assert config.system_prompt("opaque-efficient", "opaque-capable") == ( + LLM_V2_SYSTEM_PROMPT + "\n\nConfigured solver profiles:\n" + json.dumps(old_payload) + schema + ) + + +@pytest.mark.asyncio +async def test_preset_router_passes_catalog_text_and_opaque_group_names_to_judge() -> None: + catalog: Final = get_fuse_presets() + config: Final = _config(llm_v2_config=_preset_config().model_dump()) + router, client = _router(_verdict().model_dump_json(), config) + outcome: Final = await router.aclassify("Complete the supplied task") + assert outcome.tier == ComplexityTier.SIMPLE + prompt: Final = client.acompletion.call_args.kwargs["messages"][0]["content"] + payload: Final = json.loads(prompt.split("Configured solver profiles:\n")[1]) + assert payload == { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": catalog.harnesses[-1].text, + "efficient": {"model": "efficient", "profile": catalog.models[0].text}, + "capable": {"model": "capable", "profile": catalog.models[-1].text}, + } + + @pytest.mark.asyncio async def test_one_judge_fuses_whole_task_and_keeps_caller_text_out_of_system_prompt() -> None: router, client = _router(_verdict().model_dump_json()) diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 3dcb8d5af94..7d59a0590f2 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,7 +1,10 @@ from collections.abc import Mapping +from typing import Final import pytest +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets + from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, @@ -171,6 +174,52 @@ def test_validate_accepts_loadable_complexity_config(complexity_router_config): assert validate_complexity_router_config_write(complexity_router_config=complexity_router_config) is None +def _fuse_write_config(profiles: Mapping[str, object]) -> Mapping[str, object]: + return { + "classifier_type": "llm_v2", + "classifier_llm_config": {"model": "judge"}, + "tiers": {"SIMPLE": ["opaque-efficient"], "REASONING": ["opaque-capable"]}, + "llm_v2_config": {"max_quality_gap": 0.05, **profiles}, + } + + +def test_fuse_write_accepts_presets_and_custom_text_with_the_same_entitlement() -> None: + catalog: Final = get_fuse_presets() + presets: Final = _fuse_write_config( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[0].id, + } + ) + custom: Final = _fuse_write_config( + { + "efficient_profile": catalog.models[0].text, + "capable_profile": catalog.models[-1].text, + "harness": catalog.harnesses[0].text, + } + ) + assert validate_complexity_router_config_write(presets) is None + assert validate_complexity_router_config_write(custom) is None + assert claimed_capability(presets) is claimed_capability(custom) + assert claimed_capability(presets) is not None + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_fuse_write_rejects_unknown_preset_even_with_custom_text(field: str) -> None: + config: Final = _fuse_write_config( + { + "efficient_profile": "Custom efficient solver", + "capable_profile": "Custom capable solver", + "harness": "Custom runtime", + f"{field}_preset": "unknown-v1", + } + ) + violation: Final = validate_complexity_router_config_write(config) + assert violation is not None + assert f"{field}_preset" in violation + + def test_naming_check_ignores_the_config_entirely(): """The naming contract and the config's contents are separate questions with separate owners; a write may carry a config without naming a model, so neither can stand in for the other.""" diff --git a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx index 4a574ac736d..f9b0edf9508 100644 --- a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; -import { describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import userEvent from "@testing-library/user-event"; -import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs"; import ForecastClassifierConfig from "./ForecastClassifierConfig"; @@ -37,6 +37,53 @@ const fuseInitial: ComplexityRouterConfigValue = { }, }; const options = ["judge", "efficient", "capable"].map((model) => ({ value: model, label: model })); +const catalog = { + version: "catalog-v1", + models: [ + { + id: "efficient-v1", + label: "Efficient preset", + text: "Maintained efficient profile", + sources: ["https://example.com/efficient"], + model: "efficient-model", + }, + { + id: "capable-v1", + label: "Capable preset", + text: "Maintained capable profile", + sources: ["https://example.com/capable"], + model: "capable-model", + }, + ], + harnesses: [ + { + id: "runtime-v1", + label: "Runtime preset", + text: "Maintained runtime profile", + sources: ["https://example.com/runtime"], + }, + ], +}; +const presetConfig = { + efficient_profile_preset: catalog.models[0].id, + capable_profile_preset: catalog.models[1].id, + harness_preset: catalog.harnesses[0].id, + max_quality_gap: 0.05, +}; +const presetInitial = { ...fuseInitial, llm_v2_config: presetConfig }; + +beforeEach(() => { + testQueryClient.clear(); + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation(async () => Response.json(catalog)), + ); +}); + +afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); +}); function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) { const [value, setValue] = useState(initialValue); @@ -72,6 +119,118 @@ function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfi } describe("forecast classifier form", () => { + it("selects all three maintained presets, previews provenance, and saves only references", async () => { + const user = userEvent.setup(); + renderWithProviders(
); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(await screen.findByRole("option", { name: /^Efficient preset/ })); + await user.click(screen.getByRole("combobox", { name: "Capable solver profile preset" })); + await user.click(screen.getByRole("option", { name: /^Capable preset/ })); + await user.click(screen.getByRole("combobox", { name: "Harness and budget preset" })); + await user.click(screen.getByRole("option", { name: /^Runtime preset/ })); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(catalog.models[0].text); + expect(screen.getByLabelText("Efficient solver profile")).toHaveAttribute("readonly"); + expect(screen.getByLabelText("Capable solver profile")).toHaveValue(catalog.models[1].text); + expect(screen.getByLabelText("Harness and budget")).toHaveValue(catalog.harnesses[0].text); + expect(screen.getAllByText(`Catalog version: ${catalog.version}`)).toHaveLength(3); + expect(screen.getByText(`Model: ${catalog.models[0].model}`)).toBeInTheDocument(); + expect(screen.getAllByRole("link", { name: "Source 1" }).map((link) => link.getAttribute("href"))).toEqual([ + catalog.models[0].sources[0], + catalog.models[1].sources[0], + catalog.harnesses[0].sources[0], + ]); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + presetConfig, + ); + expect(fetch).toHaveBeenCalledTimes(1); + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ url: expect.stringMatching(/\/public\/complexity_router\/fuse_presets$/) }), + ); + }); + + it.each([undefined, null, "Explicit override"])( + "copies effective text to Custom and clears only that reference, override=%s", + async (override) => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + const effectiveText = override ?? catalog.models[0].text; + await waitFor(() => expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText)); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(screen.getByRole("option", { name: "Custom", exact: true })); + expect(screen.getByLabelText("Efficient solver profile")).not.toHaveAttribute("readonly"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Custom budget" } }); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + const { efficient_profile_preset: _preset, ...rest } = presetConfig; + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ + ...rest, + efficient_profile: "Custom budget", + }); + }, + ); + + it.each([ + { ...fuseInitial.llm_v2_config!, efficient_profile: catalog.models[0].text }, + { + ...presetConfig, + efficient_profile: "Explicit override", + capable_profile: "Capable override", + harness: "Harness override", + }, + ])("keeps existing custom ownership and references on an unchanged save: %j", async (settings) => { + renderWithProviders(); + await waitFor(() => expect(screen.queryByText(/Loading profile presets/)).not.toBeInTheDocument()); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("Custom"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(settings.efficient_profile); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + + it.each([true, false])( + "keeps edits and stored IDs while the pending catalog settles, success=%s", + async (success) => { + let resolveCatalog: (response: Response) => void = () => {}; + vi.mocked(fetch).mockReturnValue( + new Promise((resolve) => { + resolveCatalog = resolve; + }), + ); + const settings = { ...presetConfig, efficient_profile: "Original override" }; + renderWithProviders(); + expect(screen.getByText(/Loading profile presets/)).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Typed while loading" } }); + await act(async () => + resolveCatalog(success ? Response.json(catalog) : Response.json({ error: "unavailable" }, { status: 503 })), + ); + if (success) await screen.findAllByText(`Catalog version: ${catalog.version}`); + else expect(await screen.findByText(/Profile presets could not be loaded/)).toBeInTheDocument(); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue("Typed while loading"); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ ...settings, efficient_profile: "Typed while loading" }); + }, + ); + + it("keeps unknown saved IDs visible with unavailable previews rather than replacing them", async () => { + const settings = { ...presetConfig, efficient_profile_preset: "unavailable-v8" }; + renderWithProviders(); + await screen.findAllByText(`Catalog version: ${catalog.version}`); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("unavailable-v8"); + expect(screen.getByText("Preset preview unavailable. The saved reference is preserved")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + it("switches a populated standard router to Capability without saving hidden pools or their overrides", () => { renderWithProviders( ) : ( <> - {(["efficient_profile", "capable_profile", "harness"] as const).map((field) => { - const label = { - efficient_profile: "Efficient solver profile", - capable_profile: "Capable solver profile", - harness: "Harness and budget", - }[field]; - return ( -
- -