From 2286bf3eca414cc24e0a03b008a7a4e6b9647c44 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 18:30:16 +0900 Subject: [PATCH 001/162] 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/162] 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/162] 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/162] 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/162] 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/162] 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/162] 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/162] 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/162] 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/162] 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/162] 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 0f3c4ccfbba894c0a66d92ed44ca18f312bf75f0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 11 Sep 2026 00:00:17 +0000 Subject: [PATCH 012/162] 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 e254377049ca6f7087693cff4b99b291c9ba6010 Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:07:57 +0100 Subject: [PATCH 013/162] 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 014/162] 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 015/162] fix(azure): avoid mutable request mapping DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 7cb50ee5348..424422612db 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,11 +280,13 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) - request_params: Final = { - key: value - for key, value in optional_params.items() - if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") - } + request_params: Final = MappingProxyType( + { + key: value + for key, value in optional_params.items() + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") + } + ) return { "model": model, "messages": azure_messages, From 2fa115db2bac68ac90b0b900dc9eecb728c3bd4d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 17 Sep 2026 12:53:15 -0400 Subject: [PATCH 016/162] 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 ( -
- -