From 2286bf3eca414cc24e0a03b008a7a4e6b9647c44 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 18:30:16 +0900 Subject: [PATCH 001/463] 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 002/463] 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 003/463] 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 004/463] 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 005/463] 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 006/463] 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 007/463] 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 008/463] 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 009/463] 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 010/463] 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 011/463] 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 012/463] 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 013/463] 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 0f3c4ccfbba894c0a66d92ed44ca18f312bf75f0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 11 Sep 2026 00:00:17 +0000 Subject: [PATCH 014/463] 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 015/463] 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 016/463] 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 017/463] 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 018/463] 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 019/463] 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 61ac4f57394e022382459235baa598ac89ce00d1 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:03:56 +0000 Subject: [PATCH 020/463] 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 021/463] 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 022/463] 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 023/463] 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 024/463] 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 025/463] 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 026/463] 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 027/463] 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 028/463] 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 029/463] 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 030/463] 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 031/463] 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 032/463] 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 033/463] 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 034/463] 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 035/463] 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 036/463] 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 037/463] 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 038/463] 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 039/463] 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 040/463] 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 48bde68781c20df4d98915cc970eb65b23e343fe Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:47:14 +0000 Subject: [PATCH 041/463] fix(key_generate): use user's budget for UI session personal keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 13 ++-- .../test_key_management_endpoints.py | 77 +++++++++++++++++++ 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 802a7c3e469..4f4e56d7418 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1233,11 +1233,10 @@ async def _common_key_generation_helper( # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. - is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None - # Session tokens (lite login) carry max_budget=None to avoid a per-session - # LLM spend cap, but that None must not be read as "unlimited delegation - # authority". A personal key (no team) has no team-budget enforcement at - # request time, so a session token cannot delegate any budget for one. + # Session tokens (lite login) use their session max_budget for team keys, but + # personal keys are capped by user_max_budget when it is available. + is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID + is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None if ( user_api_key_dict.is_session_token and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -1255,7 +1254,9 @@ async def _common_key_generation_helper( }, ) delegation_ceiling: Final = ( - user_api_key_dict.max_budget + user_api_key_dict.user_max_budget + if is_ui_session_token and user_api_key_dict.user_max_budget is not None + else user_api_key_dict.max_budget if user_api_key_dict.max_budget is not None else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None) ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index cc0a7631b59..809c4183e13 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15509,6 +15509,83 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped(): assert "cannot exceed" in msg.lower() +@pytest.mark.asyncio +async def test_ui_session_token_personal_key_ceiling_is_user_budget(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=100) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-ui-session", + user_id="user-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + max_budget=1.0, + user_max_budget=500.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly + patch( # test-quality-ok: helper has no dependency injection seam for key persistence + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"} + try: + await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + team_table=None, + ) + except (HTTPException, ProxyException) as err: + msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) + assert "cannot exceed" not in msg.lower() + + +@pytest.mark.asyncio +async def test_ui_session_token_personal_key_above_user_budget_rejected(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=600) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-ui-session", + user_id="user-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + max_budget=1.0, + user_max_budget=500.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly + patch( # test-quality-ok: helper has no dependency injection seam for key persistence + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"} + with pytest.raises((HTTPException, ProxyException)) as exc_info: + await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + team_table=None, + ) + err = exc_info.value + code = getattr(err, "status_code", None) or getattr(err, "code", None) + msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) + assert str(code) == "400" + assert "cannot exceed" in msg.lower() + assert "500.0" in msg + + @pytest.mark.asyncio async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(): """ From 2fa115db2bac68ac90b0b900dc9eecb728c3bd4d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 17 Sep 2026 12:53:15 -0400 Subject: [PATCH 042/463] 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 ( -
- -