From 317bb1ea4d86a90dbaac47b77fe16129fdc1fe68 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 24 Jul 2026 01:34:10 +0000 Subject: [PATCH 01/13] fix(streaming): guard empty-choices chunks in Responses bridge and Anthropic adapter iterators Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../adapters/streaming_iterator.py | 6 +- .../adapters/transformation.py | 2 +- .../streaming_iterator.py | 6 ++ .../test_streaming_iterator_combined_chunk.py | 54 +++++++++++ .../test_empty_choices_streaming_iterator.py | 91 +++++++++++++++++++ 5 files changed, 156 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index f02333c34c8..367f2585a9b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -423,7 +423,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): will_merge_into_held = ( self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None ) - is_final_chunk = chunk.choices[0].finish_reason is not None + is_final_chunk = bool(chunk.choices) and chunk.choices[0].finish_reason is not None processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic( response=chunk, current_content_block_index=self.current_content_block_index, @@ -646,7 +646,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): will_merge_into_held = ( self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None ) - is_final_chunk = chunk.choices[0].finish_reason is not None + is_final_chunk = bool(chunk.choices) and chunk.choices[0].finish_reason is not None processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic( response=chunk, current_content_block_index=self.current_content_block_index, @@ -889,6 +889,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Example logic - customize based on your needs: # If chunk indicates a tool call + if not chunk.choices: + return False if chunk.choices[0].finish_reason is not None: return False diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 4b6617fbeac..91efc39670e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1476,7 +1476,7 @@ class LiteLLMAnthropicMessagesAdapter: applied_edits: Optional[List[AppliedEdit]] = None, ) -> Union[ContentBlockDelta, MessageBlockDelta]: ## base case - final chunk w/ finish reason - if response.choices[0].finish_reason is not None: + if response.choices and response.choices[0].finish_reason is not None: delta = MessageDelta( stop_reason=self._translate_openai_finish_reason_to_anthropic(response.choices[0].finish_reason), ) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index cf69654d15d..439c4715506 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -125,6 +125,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return None def _is_reasoning_end(self, chunk): + if not chunk.choices: + return False delta = chunk.choices[0].delta # if this indicates reasoning content, don't consider reasoning ended @@ -722,6 +724,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Change: Never return a value, just enqueue output item events if self.sent_output_item_added_event: return + if not chunk.choices: + return delta = chunk.choices[0].delta self._sequence_number += 1 @@ -1033,6 +1037,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): It's unclear how users expect litellm to translate multiple-choices-per-chunk to the responses API output. """ + if not choices: + return "" choice = choices[0] chat_completion_delta: ChatCompletionDelta = choice.delta return chat_completion_delta.content or "" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index 6973340101e..a6523f5a39d 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -143,6 +143,60 @@ def test_delayed_usage_chunk_preserves_cache_tokens(): assert message_delta["usage"]["cache_creation_input_tokens"] == 20 +def test_trailing_empty_choices_usage_chunk_emits_message_delta_usage(): + """Regression for LIT-4767. + + The trailing usage-only chunk an OpenAI-compatible provider sends when + ``include_usage`` is set has ``choices: []``. The adapter used to index + ``choices[0]`` unguarded (``is_final_chunk`` / ``_should_start_new_content_block``) + and crash with IndexError. It must instead merge the usage into the held + stop-reason chunk so ``message_delta`` still reports it. + """ + chunks = [ + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="Two."), finish_reason=None)], + ), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + ), + ModelResponseStream( + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o") + events = list(wrapper) + + message_delta = next(event for event in events if event.get("type") == "message_delta") + assert message_delta["usage"]["input_tokens"] == 10 + assert message_delta["usage"]["output_tokens"] == 5 + + +def test_leading_empty_choices_chunk_does_not_crash_stream(): + """Azure emits a leading ``prompt_filter_results`` chunk with ``choices: []`` + before any content. It must be tolerated and the following content emitted.""" + chunks = [ + ModelResponseStream(choices=[]), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="Hi"), finish_reason=None)], + ), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4), + ), + ] + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in chunks: + yield chunk + + wrapper = AnthropicStreamWrapper(completion_stream=_aiter(), model="gpt-4o") + sse = _collect_async(wrapper) + + assert "Hi" in sse + assert "message_stop" in sse + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py b/tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py new file mode 100644 index 00000000000..b7a3611501d --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py @@ -0,0 +1,91 @@ +""" +Regression tests for LIT-4767. + +When an upstream OpenAI-compatible provider emits a chunk with ``choices: []`` +(the trailing usage-only chunk every provider sends when ``include_usage`` is +set, or Azure's leading ``prompt_filter_results`` chunk), the Responses bridge +iterator used to index ``choices[0]`` unguarded and die with +``IndexError: list index out of range``, killing the whole stream. + +The empty-choices chunk must be tolerated without crashing, and the usage it +carries must still reach ``response.completed``. +""" + +from unittest.mock import AsyncMock + +from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, +) +from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +def _iterator() -> LiteLLMCompletionStreamingIterator: + return LiteLLMCompletionStreamingIterator( + model="gpt-4o", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="hi", + responses_api_request={}, + custom_llm_provider="openai", + ) + + +def _empty_choices_usage_chunk() -> ModelResponseStream: + chunk = ModelResponseStream(id="chunk-usage", model="gpt-4o", choices=[]) + chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + return chunk + + +def test_ensure_output_item_for_empty_choices_chunk_does_not_crash(): + """First chunk with no choices must not raise (traceback frame in the ticket).""" + iterator = _iterator() + # Would raise IndexError before the fix. + assert iterator._ensure_output_item_for_chunk(_empty_choices_usage_chunk()) is None + assert iterator.sent_output_item_added_event is False + + +def test_transform_empty_choices_chunk_returns_no_delta(): + """The mid/trailing usage chunk flows through transform without crashing.""" + iterator = _iterator() + # Would raise IndexError in _get_delta_string_from_streaming_choices before the fix. + assert iterator._transform_chat_completion_chunk_to_response_api_chunk(_empty_choices_usage_chunk()) is None + + +def test_is_reasoning_end_false_for_empty_choices_chunk(): + iterator = _iterator() + assert iterator._is_reasoning_end(_empty_choices_usage_chunk()) is False + + +def test_empty_choices_usage_chunk_still_reaches_response_completed(): + """End-to-end: a text chunk followed by a choices=[] usage chunk must emit + response.completed carrying the usage rather than dying mid-stream.""" + + class _SyncWrapper: + def __init__(self, chunks): + self._it = iter(chunks) + self.logging_obj = None + self.stream_options = {"include_usage": True} + + def __next__(self): + return next(self._it) + + text_chunk = ModelResponseStream( + id="chunk-1", + model="gpt-4o", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="Hi"), finish_reason=None)], + ) + iterator = _iterator() + iterator.litellm_logging_obj = None + iterator.litellm_custom_stream_wrapper = _SyncWrapper([text_chunk, _empty_choices_usage_chunk()]) + + events = list(iterator) + + completed = [e for e in events if getattr(e, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] + assert len(completed) == 1 + assert completed[0].response.usage is not None + assert completed[0].response.usage.total_tokens == 15 From 69e9edb91a9bc8efcadf876917f11cce1b31eae6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:31:38 -0700 Subject: [PATCH 02/13] test(responses): fold the empty-choices regression tests into the mapped streaming iterator test file --- .../test_empty_choices_streaming_iterator.py | 91 ------------------- .../test_streaming_iterator_transformation.py | 43 +++++++++ 2 files changed, 43 insertions(+), 91 deletions(-) delete mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py b/tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py deleted file mode 100644 index b7a3611501d..00000000000 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_empty_choices_streaming_iterator.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -Regression tests for LIT-4767. - -When an upstream OpenAI-compatible provider emits a chunk with ``choices: []`` -(the trailing usage-only chunk every provider sends when ``include_usage`` is -set, or Azure's leading ``prompt_filter_results`` chunk), the Responses bridge -iterator used to index ``choices[0]`` unguarded and die with -``IndexError: list index out of range``, killing the whole stream. - -The empty-choices chunk must be tolerated without crashing, and the usage it -carries must still reach ``response.completed``. -""" - -from unittest.mock import AsyncMock - -from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, -) -from litellm.types.llms.openai import ResponsesAPIStreamEvents -from litellm.types.utils import ( - Delta, - ModelResponseStream, - StreamingChoices, - Usage, -) - - -def _iterator() -> LiteLLMCompletionStreamingIterator: - return LiteLLMCompletionStreamingIterator( - model="gpt-4o", - litellm_custom_stream_wrapper=AsyncMock(), - request_input="hi", - responses_api_request={}, - custom_llm_provider="openai", - ) - - -def _empty_choices_usage_chunk() -> ModelResponseStream: - chunk = ModelResponseStream(id="chunk-usage", model="gpt-4o", choices=[]) - chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) - return chunk - - -def test_ensure_output_item_for_empty_choices_chunk_does_not_crash(): - """First chunk with no choices must not raise (traceback frame in the ticket).""" - iterator = _iterator() - # Would raise IndexError before the fix. - assert iterator._ensure_output_item_for_chunk(_empty_choices_usage_chunk()) is None - assert iterator.sent_output_item_added_event is False - - -def test_transform_empty_choices_chunk_returns_no_delta(): - """The mid/trailing usage chunk flows through transform without crashing.""" - iterator = _iterator() - # Would raise IndexError in _get_delta_string_from_streaming_choices before the fix. - assert iterator._transform_chat_completion_chunk_to_response_api_chunk(_empty_choices_usage_chunk()) is None - - -def test_is_reasoning_end_false_for_empty_choices_chunk(): - iterator = _iterator() - assert iterator._is_reasoning_end(_empty_choices_usage_chunk()) is False - - -def test_empty_choices_usage_chunk_still_reaches_response_completed(): - """End-to-end: a text chunk followed by a choices=[] usage chunk must emit - response.completed carrying the usage rather than dying mid-stream.""" - - class _SyncWrapper: - def __init__(self, chunks): - self._it = iter(chunks) - self.logging_obj = None - self.stream_options = {"include_usage": True} - - def __next__(self): - return next(self._it) - - text_chunk = ModelResponseStream( - id="chunk-1", - model="gpt-4o", - choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="Hi"), finish_reason=None)], - ) - iterator = _iterator() - iterator.litellm_logging_obj = None - iterator.litellm_custom_stream_wrapper = _SyncWrapper([text_chunk, _empty_choices_usage_chunk()]) - - events = list(iterator) - - completed = [e for e in events if getattr(e, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] - assert len(completed) == 1 - assert completed[0].response.usage is not None - assert completed[0].response.usage.total_tokens == 15 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 5d97b0531d6..343fc873fa4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -752,6 +752,49 @@ def test_completed_event_restores_usage_hidden_by_stream_options_none(): assert completed.response.usage.output_tokens == 5 +def _empty_choices_chunk(usage: Usage | None = None) -> ModelResponseStream: + return ModelResponseStream(id=CHAT_COMPLETION_ID, model="claude-haiku-4-5", choices=[], usage=usage) + + +@pytest.mark.asyncio +async def test_leading_empty_choices_chunk_does_not_kill_the_stream(): + """ + Azure leads some streams with a `prompt_filter_results` chunk whose `choices` is empty. + The bridge used to index `choices[0]` on it and die before the first token. + """ + iterator = _build_iterator([_empty_choices_chunk(), _chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + event_types = [getattr(event, "type", None) for event in events] + assert event_types.count(ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) == 1 + assert "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + assert event_types[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + + +@pytest.mark.asyncio +async def test_trailing_empty_choices_usage_chunk_reaches_response_completed(): + """ + With `stream_options.include_usage` (which the bridge always sets) the last upstream chunk + carries only usage and an empty `choices`. It must not crash the stream, and its usage must + still land on `response.completed`. + """ + usage: Final = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + iterator = _build_iterator([_chunk("Hello"), _chunk("", finish_reason="stop"), _empty_choices_chunk(usage)]) + + events = [event async for event in iterator] + + completed = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert completed.response.usage.input_tokens == 10 + assert completed.response.usage.output_tokens == 5 + + +def test_is_reasoning_end_ignores_empty_choices_chunk(): + assert _build_iterator([])._is_reasoning_end(_empty_choices_chunk()) is False + + def test_object_tool_call_arguments_stream_as_valid_json(): """A provider that sends decoded object arguments must still stream valid JSON. From e62ab285614d934542f8c44f7917030a44ecbbd7 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 16 Sep 2026 00:26:34 +0000 Subject: [PATCH 03/13] chore(codeowners): add ryan and kerry as owners of the cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/CODEOWNERS | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index cfa0390e836..70a50d7f06e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -4,7 +4,7 @@ /ui/nginx.conf /ui/litellm-dashboard/src/lib/http/schema.d.ts /ui/litellm-dashboard/tsconfig.tsbuildinfo -/model_prices_and_context_window.json @mateo-berri -/litellm/model_prices_and_context_window_backup.json @mateo-berri +/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri +/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri /litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri /.github/CODEOWNERS @yuneng-berri From f8fb31db3a33017dbb9a7386fcc50cb222133f44 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:51:09 +0000 Subject: [PATCH 04/13] chore(prices): sync Azure prices: 2 models azure_ai/grok-4.3: input_cost_per_token_above_200k_tokens, output_cost_per_token_above_200k_tokens, cache_read_input_token_cost_above_200k_tokens azure_ai/grok-4.6: input_cost_per_token_above_200k_tokens, output_cost_per_token_above_200k_tokens, cache_read_input_token_cost_above_200k_tokens --- litellm/model_prices_and_context_window_backup.json | 6 ++++++ model_prices_and_context_window.json | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cd458666539..88850135e19 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11313,13 +11313,16 @@ }, "azure_ai/grok-4.3": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, @@ -11331,13 +11334,16 @@ }, "azure_ai/grok-4.6": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cd458666539..88850135e19 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11313,13 +11313,16 @@ }, "azure_ai/grok-4.3": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, @@ -11331,13 +11334,16 @@ }, "azure_ai/grok-4.6": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, From 033aa8ba6d6a1d0de5fc1e16b466fcf9f329abfb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:50:36 -0700 Subject: [PATCH 05/13] fix(bedrock): forward userContext in Knowledge Base Retrieve requests The Bedrock vector store search only lifted retrievalConfiguration out of extra_body, so the caller's userContext (the Retrieve API's ACL identity) never reached Bedrock and ACL-enabled data sources answered with zero results. The transform now forwards userContext, taken from extra_body first and then from the top-level params where the OpenAI SDK's extra_body merge lands, as the caller sent it. --- .../bedrock/vector_stores/transformation.py | 21 ++++++++++ .../integrations/rag/bedrock_knowledgebase.py | 7 +++- ...est_bedrock_vector_store_transformation.py | 42 +++++++++++++++++++ tests/test_litellm/vector_stores/test_main.py | 25 +++++++++++ 4 files changed, 94 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 27c90c9d71e..ba8ce7e5625 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from copy import deepcopy from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse @@ -14,6 +15,7 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBResponse, BedrockKBRetrievalConfiguration, BedrockKBRetrievalQuery, + BedrockKBUserContext, ) from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -242,10 +244,29 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters if retrieval_config: request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config) + user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params) + if user_context is not None: + request_body["userContext"] = user_context litellm_logging_obj.model_call_details["query"] = query return url, request_body + @staticmethod + def _user_context( + extra_body: Mapping[str, object] | None, litellm_params: Mapping[str, object] + ) -> BedrockKBUserContext | None: + sources: Final = tuple(source for source in (extra_body, litellm_params) if isinstance(source, Mapping)) + found: Final = next( + ( + source[key] + for source in sources + for key in ("userContext", "user_context") + if source.get(key) is not None + ), + None, + ) + return None if found is None else cast(BedrockKBUserContext, found) + def sign_request( self, headers: dict, diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index e3aba85ed9b..7156d8101e1 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,6 +1,6 @@ from typing import Any, Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class BedrockKBLocation(TypedDict, total=False): @@ -127,6 +127,10 @@ class BedrockKBGuardrailConfiguration(TypedDict, total=False): guardrailVersion: str | None +class BedrockKBUserContext(TypedDict): + userId: ReadOnly[str] + + class BedrockKBRequest(TypedDict, total=False): """Complete request structure for Bedrock Knowledge Base retrieval.""" @@ -134,6 +138,7 @@ class BedrockKBRequest(TypedDict, total=False): nextToken: str | None retrievalConfiguration: BedrockKBRetrievalConfiguration | None retrievalQuery: BedrockKBRetrievalQuery + userContext: ReadOnly[BedrockKBUserContext | None] ######################################################################### diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 7b04efa17dc..e435f0f7a8b 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -82,6 +82,7 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): == "HYBRID" ) assert "unrelatedField" not in body + assert "userContext" not in body def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results(): @@ -152,3 +153,44 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body() ]["value"] == "a" ) + + +def _search_body(extra_body, litellm_params): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"max_num_results": 3}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params=litellm_params, + extra_body=extra_body, + ) + return body + + +def test_transform_search_request_forwards_user_context_from_extra_body(): + body = _search_body(extra_body={"userContext": {"userId": "alice@example.com"}}, litellm_params={}) + + assert body["userContext"] == {"userId": "alice@example.com"} + assert body["retrievalConfiguration"] == {"vectorSearchConfiguration": {"numberOfResults": 3}} + + +def test_transform_search_request_forwards_top_level_user_context_from_litellm_params(): + body = _search_body( + extra_body=None, + litellm_params={"vector_store_id": "kb123", "user_context": {"userId": "bob@example.com"}}, + ) + + assert body["userContext"] == {"userId": "bob@example.com"} + + +def test_transform_search_request_prefers_extra_body_user_context_over_top_level(): + body = _search_body( + extra_body={"userContext": {"userId": "alice@example.com"}}, + litellm_params={"userContext": {"userId": "bob@example.com"}}, + ) + + assert body["userContext"] == {"userId": "alice@example.com"} diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/test_litellm/vector_stores/test_main.py index e3575c33b17..1c968126c42 100644 --- a/tests/test_litellm/vector_stores/test_main.py +++ b/tests/test_litellm/vector_stores/test_main.py @@ -7,6 +7,7 @@ executor, and it must never leak into litellm_params/kwargs where logging would model_dump() it (the #19550 serialization trap). """ +import json from unittest.mock import MagicMock, patch import pytest @@ -15,6 +16,7 @@ import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.vector_stores.main import search MOCK_SEARCH_RESPONSE = { @@ -89,3 +91,26 @@ def test_search_router_not_in_litellm_params(): litellm_params = mock_handler.call_args.kwargs["litellm_params"] assert "router" not in litellm_params.model_dump(exclude_none=True) assert getattr(litellm_params, "router", None) is None + + +def test_search_forwards_top_level_user_context_to_bedrock_retrieve(): + """Regression (LIT-4415): a top-level userContext, the shape the OpenAI SDK's extra_body + produces on the proxy path, reaches the Bedrock Retrieve request body.""" + client = MagicMock(spec=HTTPHandler) + client.post.return_value = MagicMock(status_code=200, json=MagicMock(return_value={"retrievalResults": []})) + + search( + vector_store_id="kb123", + query="q", + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + aws_access_key_id="test-key-id", + aws_secret_access_key="test-secret-key", + userContext={"userId": "alice@example.com"}, + client=client, + litellm_logging_obj=MagicMock(), + ) + + posted = json.loads(client.post.call_args.kwargs["data"]) + assert posted["userContext"] == {"userId": "alice@example.com"} + assert posted["retrievalQuery"] == {"text": "q"} From b893e6b926dde5136a2b733fa5ade4d038ed120c Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:06:19 +0000 Subject: [PATCH 06/13] chore(prices): sync Google Gemini prices: 22 models gemini/gemini-2.5-flash: max_tokens, max_output_tokens, supports_audio_input gemini/gemini-2.5-flash-image: max_input_tokens, supports_web_search, supports_audio_input, supports_response_schema, supports_function_calling gemini/gemini-2.5-flash-lite: max_tokens, max_output_tokens, supports_audio_input gemini-2.5-flash-native-audio-preview-12-2025: supports_vision, max_input_tokens, supports_web_search, supports_response_schema, supports_function_calling gemini/gemini-2.5-flash-native-audio-preview-12-2025: supports_vision, max_input_tokens, supports_web_search, supports_response_schema, supports_function_calling gemini-2.5-flash-preview-tts: max_tokens, max_input_tokens, max_output_tokens, supports_web_search, supports_audio_input, supports_response_schema, supports_function_calling gemini/gemini-2.5-flash-preview-tts: max_tokens, max_input_tokens, max_output_tokens, supports_web_search, supports_audio_input, supports_response_schema, supports_function_calling gemini/gemini-2.5-pro: max_tokens, max_output_tokens gemini/gemini-2.5-pro-preview-tts: max_tokens, supports_vision, max_input_tokens, max_output_tokens, supports_web_search, supports_audio_input, supports_response_schema, supports_function_calling gemini/gemini-3-flash-preview: max_tokens, max_output_tokens, supports_audio_input gemini/gemini-3-pro-image: supports_response_schema gemini/gemini-3.1-flash-image: max_input_tokens, supports_response_schema gemini/gemini-3.1-flash-lite-image: supports_web_search, supports_function_calling gemini-3.1-flash-live-preview: supports_response_schema gemini/gemini-3.1-flash-live-preview: supports_response_schema gemini/gemini-3.1-flash-tts-preview: supports_web_search, supports_response_schema, supports_function_calling gemini/gemini-3.5-flash: max_tokens, max_output_tokens gemini/gemini-3.5-live-translate-preview: supports_web_search, supports_response_schema, supports_function_calling gemini/gemini-3.5-transcribe: supports_function_calling gemini/gemini-3.5-transcribe-live: supports_function_calling gemini/gemini-embedding-2: supports_vision, supports_audio_input gemini/gemini-omni-1.1-flash: max_input_tokens --- ...odel_prices_and_context_window_backup.json | 118 ++++++++++++------ model_prices_and_context_window.json | 118 ++++++++++++------ 2 files changed, 158 insertions(+), 78 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 88850135e19..dd21bbf0b25 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26168,7 +26168,9 @@ "output_vector_size": 3072, "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supports_audio_input": true, "supports_multimodal": true, + "supports_vision": true, "tpm": 10000000 }, "gemini/gemini-1.5-flash": { @@ -26309,8 +26311,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26356,6 +26358,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-image": { @@ -26368,7 +26371,7 @@ "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "gemini", "supports_reasoning": false, - "max_input_tokens": 32768, + "max_input_tokens": 65536, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26394,22 +26397,23 @@ "image" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": true, + "supports_web_search": false, "tpm": 8000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "supports_audio_input": false, "supports_image_size": false }, "gemini/gemini-3-pro-image": { @@ -26447,7 +26451,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26550,7 +26554,7 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", - "max_input_tokens": 65536, + "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26577,7 +26581,7 @@ "supports_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26659,12 +26663,13 @@ "text", "image" ], - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": false, "supports_reasoning": false, "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, + "supports_web_search": false, "tpm": 4000000 }, "gemini/deep-research-pro-preview-12-2025": { @@ -26717,8 +26722,8 @@ "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 4e-07, "output_cost_per_token": 4e-07, @@ -26764,6 +26769,7 @@ "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { @@ -27011,6 +27017,9 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, @@ -27019,7 +27028,11 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, @@ -27033,8 +27046,8 @@ "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, @@ -27344,8 +27357,8 @@ "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, @@ -27394,7 +27407,8 @@ "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "output_cost_per_token_flex": 1.5e-06, + "supports_audio_input": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -27403,8 +27417,8 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, @@ -28118,9 +28132,9 @@ "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2e-05, @@ -28133,19 +28147,20 @@ "audio" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, + "supports_vision": false, + "supports_web_search": false, "tpm": 10000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "supports_audio_input": false }, "gemini/gemini-exp-1114": { "input_cost_per_token": 0, @@ -55923,7 +55938,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -55943,7 +55958,11 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -55977,7 +55996,8 @@ "supports_vision": true, "supports_web_search": true, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, @@ -56039,7 +56059,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -56061,7 +56081,11 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini/gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -56097,7 +56121,8 @@ "tpm": 250000, "rpm": 10, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, @@ -56114,19 +56139,29 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-flash-latest": { "cache_read_input_token_cost": 3e-08, @@ -58700,6 +58735,9 @@ ], "supports_audio_input": true, "supports_audio_output": true, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false, "tpm": 250000 }, "gemini/gemini-3.5-transcribe": { @@ -58721,7 +58759,8 @@ ], "supports_audio_input": true, "tpm": 800000, - "rpm": 2000 + "rpm": 2000, + "supports_function_calling": false }, "gemini/gemini-3.5-transcribe-live": { "input_cost_per_audio_token": 3.5e-06, @@ -58741,7 +58780,8 @@ ], "supports_audio_input": true, "tpm": 250000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false }, "vertex_ai/gemini-3.5-transcribe-preview": { "input_cost_per_audio_token": 2e-06, @@ -61380,7 +61420,7 @@ "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", - "max_input_tokens": 131072, + "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 88850135e19..dd21bbf0b25 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26168,7 +26168,9 @@ "output_vector_size": 3072, "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supports_audio_input": true, "supports_multimodal": true, + "supports_vision": true, "tpm": 10000000 }, "gemini/gemini-1.5-flash": { @@ -26309,8 +26311,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26356,6 +26358,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-image": { @@ -26368,7 +26371,7 @@ "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "gemini", "supports_reasoning": false, - "max_input_tokens": 32768, + "max_input_tokens": 65536, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26394,22 +26397,23 @@ "image" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": true, + "supports_web_search": false, "tpm": 8000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "supports_audio_input": false, "supports_image_size": false }, "gemini/gemini-3-pro-image": { @@ -26447,7 +26451,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26550,7 +26554,7 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", - "max_input_tokens": 65536, + "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26577,7 +26581,7 @@ "supports_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26659,12 +26663,13 @@ "text", "image" ], - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": false, "supports_reasoning": false, "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, + "supports_web_search": false, "tpm": 4000000 }, "gemini/deep-research-pro-preview-12-2025": { @@ -26717,8 +26722,8 @@ "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 4e-07, "output_cost_per_token": 4e-07, @@ -26764,6 +26769,7 @@ "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { @@ -27011,6 +27017,9 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, @@ -27019,7 +27028,11 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, @@ -27033,8 +27046,8 @@ "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, @@ -27344,8 +27357,8 @@ "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, @@ -27394,7 +27407,8 @@ "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "output_cost_per_token_flex": 1.5e-06, + "supports_audio_input": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -27403,8 +27417,8 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, @@ -28118,9 +28132,9 @@ "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2e-05, @@ -28133,19 +28147,20 @@ "audio" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, + "supports_vision": false, + "supports_web_search": false, "tpm": 10000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "supports_audio_input": false }, "gemini/gemini-exp-1114": { "input_cost_per_token": 0, @@ -55923,7 +55938,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -55943,7 +55958,11 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -55977,7 +55996,8 @@ "supports_vision": true, "supports_web_search": true, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, @@ -56039,7 +56059,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -56061,7 +56081,11 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini/gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -56097,7 +56121,8 @@ "tpm": 250000, "rpm": 10, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, @@ -56114,19 +56139,29 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-flash-latest": { "cache_read_input_token_cost": 3e-08, @@ -58700,6 +58735,9 @@ ], "supports_audio_input": true, "supports_audio_output": true, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false, "tpm": 250000 }, "gemini/gemini-3.5-transcribe": { @@ -58721,7 +58759,8 @@ ], "supports_audio_input": true, "tpm": 800000, - "rpm": 2000 + "rpm": 2000, + "supports_function_calling": false }, "gemini/gemini-3.5-transcribe-live": { "input_cost_per_audio_token": 3.5e-06, @@ -58741,7 +58780,8 @@ ], "supports_audio_input": true, "tpm": 250000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false }, "vertex_ai/gemini-3.5-transcribe-preview": { "input_cost_per_audio_token": 2e-06, @@ -61380,7 +61420,7 @@ "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", - "max_input_tokens": 131072, + "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", From b70ddc2fd8cef1c86807974fb2cbef06f92bc85f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:14:32 +0000 Subject: [PATCH 07/13] ci(rust): split rust jobs, use nextest and Swatinem/rust-cache Split the Rust workflow into fmt, clippy, nextest and wheel jobs so they run in parallel, replace manual actions/cache with Swatinem/rust-cache, and install a pinned checksum-verified cargo-nextest. Make two python-bridge tests self-contained so they pass when nextest runs each test in its own process. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 81 +++++++++++++------ .../crates/python-bridge/src/lifecycle/mod.rs | 52 +++++++----- .../crates/python-bridge/src/marshal.rs | 1 + 3 files changed, 89 insertions(+), 45 deletions(-) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index c4847aca20d..00c6fed8ae3 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -68,9 +68,9 @@ env: CARGO_TERM_COLOR: always jobs: - rust-lint: + rust-fmt: runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: 5 defaults: run: working-directory: litellm-rust @@ -81,24 +81,67 @@ jobs: - run: rustup toolchain install --no-self-update - - run: cargo fmt --check + - run: cargo fmt --all --check - - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + rust-clippy: + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: litellm-rust + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: - path: | - ~/.cargo/registry - ~/.cargo/git - litellm-rust/target - key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} - restore-keys: | - ${{ runner.os }}-cargo-${{ github.job }}- + persist-credentials: false + + - run: rustup toolchain install --no-self-update + + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + with: + workspaces: litellm-rust + cache-on-failure: true - run: cargo clippy --workspace --all-targets --locked -- -D warnings rust-test: runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 20 + defaults: + run: + working-directory: litellm-rust + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - run: rustup toolchain install --no-self-update + + - name: Install cargo-nextest 0.9.143 + working-directory: ${{ runner.temp }} + run: | + curl -fsSL --retry 3 -o cargo-nextest.tar.gz \ + https://github.com/nextest-rs/nextest/releases/download/cargo-nextest-0.9.143/cargo-nextest-0.9.143-x86_64-unknown-linux-gnu.tar.gz + echo "66786b9abe23920d022a182d1416b1bbc8130dd4872a9553d76985a1708dcd1e cargo-nextest.tar.gz" | sha256sum -c - + mkdir -p bin + tar xzf cargo-nextest.tar.gz -C bin cargo-nextest + echo "$PWD/bin" >> "$GITHUB_PATH" + + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + with: + workspaces: litellm-rust + cache-on-failure: true + + - run: cargo nextest run --workspace --locked + + - run: cargo test --workspace --doc --locked + + rust-wheel: + runs-on: ubuntu-latest + timeout-minutes: 30 steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: @@ -114,18 +157,10 @@ jobs: - run: rustup toolchain install --no-self-update - - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: - path: | - ~/.cargo/registry - ~/.cargo/git - litellm-rust/target - key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} - restore-keys: | - ${{ runner.os }}-cargo-${{ github.job }}- - - - run: cargo test --workspace --locked - working-directory: litellm-rust + workspaces: litellm-rust + cache-on-failure: true - run: uv build --wheel --out-dir dist diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index cf9f31c3c13..c4b8d8eaae0 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -599,6 +599,34 @@ mod tests { static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); + fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> { + py.run( + pyo3::ffi::c_str!( + r#" +import sys +import types + +sys.modules.setdefault('litellm', types.ModuleType('litellm')) +sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) +"# + ), + None, + None, + ) + .unwrap(); + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap() + } + fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> { py.import("litellm.litellm_core_utils.logging_worker")? .setattr("GLOBAL_LOGGING_WORKER", worker) @@ -773,17 +801,7 @@ mod tests { .unwrap_or_else(|error| error.into_inner()); Python::initialize(); Python::attach(|py| { - let source = std::ffi::CString::new(include_str!( - "../../../../../litellm/rust_bridge/lifecycle.py" - )) - .unwrap(); - PyModule::from_code( - py, - &source, - pyo3::ffi::c_str!("lifecycle.py"), - pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), - ) - .unwrap(); + install_lifecycle_module(py); let route = SyntheticRoute( PythonCallState::new( py, @@ -819,17 +837,7 @@ mod tests { Python::initialize(); Python::attach(|py| { py.import("asyncio").unwrap(); - let source = std::ffi::CString::new(include_str!( - "../../../../../litellm/rust_bridge/lifecycle.py" - )) - .unwrap(); - let module = PyModule::from_code( - py, - &source, - pyo3::ffi::c_str!("lifecycle.py"), - pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), - ) - .unwrap(); + let module = install_lifecycle_module(py); let locals = PyDict::new(py); locals .set_item("drive", module.getattr("drive").unwrap()) diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 9038eb971b3..7f00298905f 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -190,6 +190,7 @@ mod tests { #[test] fn required_shapes_preserve_nested_values_and_existing_errors() { + Python::initialize(); let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]); assert_eq!( Value::Array(required_array("messages", nested.clone()).unwrap()), From e46106e20ba5eeb1a02db8fbbc171df6d4bab211 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 19:21:50 +0000 Subject: [PATCH 08/13] test: drop gemini-3.1-flash-lite-image capability pins The per-route capability test hardcoded vendor facts, including function calling support on the gemini route, which the live model card says is not supported. Keep the backup-matches-main invariant and the routing tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...est_gemini_3_1_flash_lite_image_pricing.py | 104 ------------------ 1 file changed, 104 deletions(-) diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py index 10d1d6fecd1..250b587aaf1 100644 --- a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py @@ -3,14 +3,7 @@ from pathlib import Path import pytest -import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.types.utils import ( - ImageObject, - ImageResponse, - ImageUsage, - ImageUsageInputTokensDetails, -) REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -21,94 +14,12 @@ GEMINI = "gemini/gemini-3.1-flash-lite-image" VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) -INPUT_COST = 2.5e-07 -INPUT_COST_BATCHES = 1.25e-07 -OUTPUT_TEXT_COST = 1.5e-06 -OUTPUT_TEXT_COST_BATCHES = 7.5e-07 -OUTPUT_IMAGE_TOKEN_COST = 3e-05 -OUTPUT_COST_PER_1K_IMAGE = 0.0336 -INPUT_COST_PER_IMAGE = 0.00028 -CACHE_READ_COST = 2.5e-08 -MAX_INPUT_TOKENS = 65536 -MAX_OUTPUT_TOKENS = 4096 -TOKENS_PER_1K_IMAGE = 1120 - -SHARED_FIELDS = { - "mode": "image_generation", - "input_cost_per_token": INPUT_COST, - "input_cost_per_token_batches": INPUT_COST_BATCHES, - "input_cost_per_image": INPUT_COST_PER_IMAGE, - "output_cost_per_token": OUTPUT_TEXT_COST, - "output_cost_per_token_batches": OUTPUT_TEXT_COST_BATCHES, - "output_cost_per_image": OUTPUT_COST_PER_1K_IMAGE, - "output_cost_per_image_token": OUTPUT_IMAGE_TOKEN_COST, - "max_input_tokens": MAX_INPUT_TOKENS, - "max_output_tokens": MAX_OUTPUT_TOKENS, - "max_tokens": MAX_OUTPUT_TOKENS, - "supported_endpoints": ["/v1/chat/completions", "/v1/completions", "/v1/batch"], - "supported_output_modalities": ["text", "image"], - "supports_reasoning": False, - "supports_response_schema": False, - "supports_system_messages": True, - "supports_vision": True, -} - -VERTEX_ROUTE_FIELDS = { - "litellm_provider": "vertex_ai-language-models", - "cache_read_input_token_cost": CACHE_READ_COST, - "supported_modalities": ["text", "image", "video"], - "supports_function_calling": False, - "supports_pdf_input": True, - "supports_prompt_caching": True, - "supports_video_input": True, -} - -PER_ROUTE_FIELDS = { - UNPREFIXED: VERTEX_ROUTE_FIELDS, - VERTEX: VERTEX_ROUTE_FIELDS, - GEMINI: { - "litellm_provider": "gemini", - "supported_modalities": ["text", "image"], - "supports_function_calling": True, - "supports_prompt_caching": False, - "rpm": 1000, - "tpm": 4000000, - }, -} - -GROUNDING_FIELDS = ( - "supports_web_search", - "search_context_cost_per_query", - "web_search_billing_unit", -) - def _load(path: Path) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - - -@pytest.mark.parametrize("model", ALL_KEYS) -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_per_route_capabilities_match_model_cards(model: str, path: Path): - info = _load(path)[model] - for field, value in PER_ROUTE_FIELDS[model].items(): - assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" - - @pytest.mark.parametrize("model", ALL_KEYS) def test_backup_matches_main(model: str): assert _load(BACKUP_PATH).get(model) == _load(MAIN_PATH).get(model) @@ -124,18 +35,3 @@ def test_vertex_prefix_routes_to_vertex(): routed_model, provider, _, _ = get_llm_provider(model=VERTEX) assert routed_model == UNPREFIXED assert provider == "vertex_ai" - - -def _one_k_image_response() -> ImageResponse: - return ImageResponse( - data=[ImageObject(b64_json="img1")], - usage=ImageUsage( - input_tokens=50 + TOKENS_PER_1K_IMAGE, - input_tokens_details=ImageUsageInputTokensDetails( - text_tokens=50, - image_tokens=TOKENS_PER_1K_IMAGE, - ), - output_tokens=TOKENS_PER_1K_IMAGE, - total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, - ), - ) From e8b5632c20676b1e5ec74de5cd21ac38875d5fa6 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:22:02 +0000 Subject: [PATCH 09/13] ci(rust): run the token counter timing test alone under nextest Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/.config/nextest.toml | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 litellm-rust/.config/nextest.toml diff --git a/litellm-rust/.config/nextest.toml b/litellm-rust/.config/nextest.toml new file mode 100644 index 00000000000..1762a151573 --- /dev/null +++ b/litellm-rust/.config/nextest.toml @@ -0,0 +1,3 @@ +[[profile.default.overrides]] +filter = "test(long_repeated_runs_stay_cheap)" +threads-required = "num-cpus" From e850232f0208fb81c1187143bf9694b79d69149f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:43:42 +0000 Subject: [PATCH 10/13] test(rust): assert merge cost scales linearly instead of a wall-clock bound Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/.config/nextest.toml | 3 --- .../crates/token-counter/src/tiktoken.rs | 19 +++++++++++++------ 2 files changed, 13 insertions(+), 9 deletions(-) delete mode 100644 litellm-rust/.config/nextest.toml diff --git a/litellm-rust/.config/nextest.toml b/litellm-rust/.config/nextest.toml deleted file mode 100644 index 1762a151573..00000000000 --- a/litellm-rust/.config/nextest.toml +++ /dev/null @@ -1,3 +0,0 @@ -[[profile.default.overrides]] -filter = "test(long_repeated_runs_stay_cheap)" -threads-required = "num-cpus" diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index c479ae01be9..7a9e71ed587 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -195,14 +195,21 @@ mod tests { } #[test] - fn long_repeated_runs_stay_cheap() { + fn long_repeated_runs_cost_close_to_linear() { let ranks = ranks(); let mut scratch = MergeScratch::default(); - let piece = vec![b' '; 1 << 20]; - let started = std::time::Instant::now(); - let count = ranks.count_piece(&piece, &mut scratch); - assert!(count > 0); - assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed()); + let mut time = |len: usize| { + let piece = vec![b' '; len]; + let started = std::time::Instant::now(); + assert!(ranks.count_piece(&piece, &mut scratch) > 0); + started.elapsed() + }; + let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); + let large = time(1 << 18); + assert!( + large < small * 64, + "{small:?} for 2^14 bytes, {large:?} for 2^18" + ); } #[test] From 3f15dcd96b10206becb807a9d8bc57f2f9323096 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:47:48 +0000 Subject: [PATCH 11/13] ci(rust): fold fmt into the clippy job Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 00c6fed8ae3..e56726deb53 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -68,22 +68,7 @@ env: CARGO_TERM_COLOR: always jobs: - rust-fmt: - runs-on: ubuntu-latest - timeout-minutes: 5 - defaults: - run: - working-directory: litellm-rust - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - run: rustup toolchain install --no-self-update - - - run: cargo fmt --all --check - - rust-clippy: + rust-lint: runs-on: ubuntu-latest timeout-minutes: 15 defaults: @@ -96,6 +81,8 @@ jobs: - run: rustup toolchain install --no-self-update + - run: cargo fmt --all --check + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: workspaces: litellm-rust From 48df3d5a48d87e45538f02625b5444e4ec867d6a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:54:46 +0000 Subject: [PATCH 12/13] ci(rust): install nextest via pinned taiki-e/install-action Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index e56726deb53..551f783d4f9 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -107,15 +107,9 @@ jobs: - run: rustup toolchain install --no-self-update - - name: Install cargo-nextest 0.9.143 - working-directory: ${{ runner.temp }} - run: | - curl -fsSL --retry 3 -o cargo-nextest.tar.gz \ - https://github.com/nextest-rs/nextest/releases/download/cargo-nextest-0.9.143/cargo-nextest-0.9.143-x86_64-unknown-linux-gnu.tar.gz - echo "66786b9abe23920d022a182d1416b1bbc8130dd4872a9553d76985a1708dcd1e cargo-nextest.tar.gz" | sha256sum -c - - mkdir -p bin - tar xzf cargo-nextest.tar.gz -C bin cargo-nextest - echo "$PWD/bin" >> "$GITHUB_PATH" + - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8 + with: + tool: cargo-nextest@0.9.143 - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: From 884087f01cf902bab71930affdd686c8d42ec1c3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 12:55:29 -0700 Subject: [PATCH 13/13] test(bedrock): type the vector store search test helper --- .../test_bedrock_vector_store_transformation.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index e435f0f7a8b..ab5a2531461 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig @@ -155,9 +156,9 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body() ) -def _search_body(extra_body, litellm_params): - config = BedrockVectorStoreConfig() - mock_log = MagicMock() +def _search_body(extra_body: dict[str, object] | None, litellm_params: dict[str, object]) -> dict[str, object]: + config: Final = BedrockVectorStoreConfig() + mock_log: Final = MagicMock() mock_log.model_call_details = {} _, body = config.transform_search_vector_store_request( vector_store_id="kb123",