From ab07b7207e37e52f1ac8126817d101c90f5d01a0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:59:57 -0700 Subject: [PATCH 1/9] fix(vertex-ai): add batch above 200k and image output batch prices for gemini models (#48945) * fix(vertex-ai): add batch and flex above 200k and image output tier prices for gemini models * fix(vertex-ai): drop flex above 200k and image priority fields the cost map schema rejects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 28 ++++++++++++++++--- model_prices_and_context_window.json | 28 ++++++++++++++++--- 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dc043dc4d26..8719565a96f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -28853,6 +28853,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -28910,6 +28911,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -28963,12 +28965,13 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29014,6 +29017,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -29059,9 +29063,10 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29106,6 +29111,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -29672,10 +29678,12 @@ }, "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, @@ -29690,6 +29698,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -29701,6 +29710,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -30229,6 +30239,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -30240,6 +30251,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -51910,6 +51922,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -51968,6 +51981,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -52002,13 +52016,14 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52034,6 +52049,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -52061,10 +52077,11 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52089,6 +52106,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -74400,6 +74418,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -74445,6 +74464,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dc043dc4d26..8719565a96f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -28853,6 +28853,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -28910,6 +28911,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -28963,12 +28965,13 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29014,6 +29017,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -29059,9 +29063,10 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29106,6 +29111,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -29672,10 +29678,12 @@ }, "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, @@ -29690,6 +29698,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -29701,6 +29710,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -30229,6 +30239,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -30240,6 +30251,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -51910,6 +51922,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -51968,6 +51981,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -52002,13 +52016,14 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52034,6 +52049,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -52061,10 +52077,11 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52089,6 +52106,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -74400,6 +74418,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -74445,6 +74464,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, From d06b889d413237881e3dc27486e4345dd7aa18f9 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:00:29 -0700 Subject: [PATCH 2/9] fix(streaming): keep merged think tags balanced and split SSE lines only on CR/LF (#48755) * fix(vertex_ai): split Gemini SSE lines on CR/LF only so U+2028 in text keeps the stream The Gemini and Vertex AI streaming path read the SSE body with httpx iter_lines()/aiter_lines(). Those follow str.splitlines(), which also breaks on U+0085, U+2028, U+2029, \x0b, \x0c and \x1c-\x1e. JSON leaves U+0085, U+2028 and U+2029 unescaped inside strings, so a data: line whose text holds one of them was cut in two. The tail is not a data: line and the head never parses, and the accumulate-and-retry fallback then carries the unparsed prefix forward, so every later event is appended to it and the stream ends with no more content and no usage. SSE lines end only at CRLF, LF or CR. The new iter_sse_lines/aiter_sse_lines split decoded text on exactly those, treating a CR at the end of one chunk and an LF at the start of the next as one break. The first request and every long-decode hop now read through them. Co-Authored-By: Claude Opus 5.5 * fix(streaming): keep merged think tags balanced and split shared SSE lines only on CR/LF Close an open block on the first content or tool-call delta and at stream end, reopen it when reasoning returns, and route the shared chat completions stream through the SSE line splitter so U+2028 and U+2029 inside JSON no longer drop chunks. * fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator * fix(streaming): split every SSE reader on CR/LF only and keep the logged think block closed Switch the remaining httpx stream readers (Anthropic chat, Google GenAI, Interactions, Codestral, openai_like, Predibase, LangGraph, Vertex Agent Engine, A2A providers, Azure AI agents, sandboxes, the A2A endpoint, and others) from iter_lines/aiter_lines to iter_sse_lines/aiter_sse_lines, so a raw U+2028 or U+2029 inside a JSON string no longer splits an event. Keep the chunk that closes a think block in the logged and cached response when a stream ends without a finish_reason, close the block before a legacy function_call, and make sync Codestral FIM streaming send stream=true and use the text-completion-codestral stream parser. Test fakes that only mocked iter_lines now stream real httpx bytes and assert the content they feed in. * fix(streaming): split buffered SSE parsers only on CR and LF, shard the codestral tests Buffered SSE and JSONL parsers used str.splitlines(), which also breaks on U+2028, U+2029 and U+0085 that JSON allows raw inside strings. They now use split_sse_lines from sse_lines.py, which also gains a str.split fast path. Adds tests/unit/llms/codestral to the unit workflow shard. * fix(streaming): keep U+2028 in restamped Anthropic frames and fold CRLF before splitting SSE lines * test: type the new stream tests, drop a comment, and benchmark the SSE line splitter * test: type the iterator class parameters in the line-separator stream tests * fix(responses): stream tool call arguments opened in a text chunk * fix(responses): hold a tool call that shares a chunk with text until the text ends * test: type the a2a SSE handler helper and the sandbox fake responses --------- Co-authored-by: Tim Vink Co-authored-by: Claude Opus 5.5 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: michelligabriele --- .github/workflows/test-unit.yml | 1 + .../bedrock_agentcore/transformation.py | 5 +- .../providers/watsonx_orchestrate/handler.py | 5 +- .../transformation.py | 3 +- litellm/google_genai/streaming_iterator.py | 9 +- litellm/interactions/streaming_iterator.py | 5 +- .../litellm_core_utils/streaming_handler.py | 57 +++--- litellm/llms/anthropic/chat/handler.py | 5 +- .../anthropic/messages/stream_assembly.py | 3 +- litellm/llms/azure_ai/agents/handler.py | 3 +- litellm/llms/base_llm/base_model_iterator.py | 2 +- .../base_llm/guardrail_translation/utils.py | 3 +- .../llms/base_llm/sandbox/transformation.py | 3 +- litellm/llms/bedrock/files/transformation.py | 3 +- .../guardrail_translation/handler.py | 3 +- litellm/llms/bedrock/search/transformation.py | 5 +- .../llms/chatgpt/responses/transformation.py | 3 +- litellm/llms/codestral/completion/handler.py | 9 +- litellm/llms/custom_httpx/aiohttp_handler.py | 3 +- litellm/llms/custom_httpx/llm_http_handler.py | 5 +- litellm/llms/custom_httpx/sse_lines.py | 74 ++++++++ .../guardrail_translation/handler.py | 3 +- .../huggingface/embedding/transformation.py | 3 +- litellm/llms/langgraph/chat/sse_iterator.py | 5 +- litellm/llms/oobabooga/chat/oobabooga.py | 3 +- litellm/llms/openai_like/chat/handler.py | 9 +- litellm/llms/predibase/chat/handler.py | 5 +- .../vertex_ai/agent_engine/transformation.py | 5 +- .../vertex_and_google_ai_studio_gemini.py | 9 +- .../proxy/agent_endpoints/a2a_endpoints.py | 3 +- .../streaming_model_restamp.py | 4 +- .../anthropic_passthrough_logging_handler.py | 7 +- .../tinyfish_passthrough_logging_handler.py | 3 +- .../managed_id_rewriter.py | 3 +- .../response_polling/background_streaming.py | 3 +- .../streaming_iterator.py | 13 +- tests/benchmarks/test_sse_lines_benchmarks.py | 27 +++ .../test_bedrock_agentcore_a2a.py | 32 ++++ ...test_watsonx_orchestrate_transformation.py | 20 ++- ...responses_transformation_transformation.py | 18 ++ .../test_google_genai_streaming_iterator.py | 83 +++++---- .../test_interactions_streaming_iterator.py | 45 ++++- .../test_streaming_handler.py | 114 +++++++++++- .../chat/test_anthropic_chat_handler.py | 59 ++++++- .../messages/test_stream_assembly.py | 24 +++ .../unit/llms/azure_ai/agents/test_handler.py | 23 ++- .../azure_ai/agents/test_transformation.py | 20 +-- .../guardrail_translation/test_utils.py | 17 ++ .../test_bedrock_files_transformation.py | 14 ++ .../guardrail_translation/test_handler.py | 14 +- .../test_agentcore_search_transformation.py | 15 ++ .../test_chatgpt_responses_transformation.py | 31 ++++ .../test_cloudflare_transformation.py | 47 ++++- tests/unit/llms/codestral/__init__.py | 0 .../llms/codestral/completion/__init__.py | 0 .../test_codestral_completion_handler.py | 50 ++++++ .../llms/custom_httpx/test_aiohttp_handler.py | 37 ++++ .../unit/llms/custom_httpx/test_sse_lines.py | 64 +++++++ .../test_databricks_chat_transformation.py | 14 +- .../test_fireworks_ai_chat_transformation.py | 6 +- .../test_google_genai_guardrail_handler.py | 16 ++ .../chat/test_langgraph_sse_iterator.py | 32 ++++ .../llms/oobabooga/chat/test_oobabooga.py | 29 ++++ .../chat/test_openai_like_chat_handler.py | 37 +++- .../test_snowflake_chat_transformation.py | 30 +--- .../llms/test_predibase_transformation.py | 47 +++++ tests/unit/llms/triton/test_triton.py | 11 +- .../agent_engine/test_transformation.py | 47 ++++- ...test_vertex_and_google_ai_studio_gemini.py | 64 ++++++- .../llms/vertex_ai/test_http_status_201.py | 8 +- .../agent_endpoints/test_a2a_endpoints.py | 138 +++++++-------- .../test_streaming_model_restamp.py | 23 +++ ...t_anthropic_passthrough_logging_handler.py | 48 +++++ ...st_tinyfish_passthrough_logging_handler.py | 7 + .../test_managed_id_rewriter.py | 33 +++- .../test_response_polling_handler.py | 31 +++- .../test_streaming_iterator_transformation.py | 100 +++++++++++ tests/unit/sandbox/test_e2b_sandbox.py | 28 ++- .../unit/sandbox/test_opensandbox_sandbox.py | 13 +- tests/unit/test_groq_streaming_encoding.py | 164 ++++-------------- 80 files changed, 1554 insertions(+), 413 deletions(-) create mode 100644 litellm/llms/custom_httpx/sse_lines.py create mode 100644 tests/benchmarks/test_sse_lines_benchmarks.py create mode 100644 tests/unit/llms/codestral/__init__.py create mode 100644 tests/unit/llms/codestral/completion/__init__.py create mode 100644 tests/unit/llms/codestral/completion/test_codestral_completion_handler.py create mode 100644 tests/unit/llms/custom_httpx/test_sse_lines.py create mode 100644 tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2ddf21bca17..80d5018fe9a 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -245,6 +245,7 @@ jobs: tests/unit/llms/clarifai tests/unit/llms/claude_code tests/unit/llms/cloudflare + tests/unit/llms/codestral tests/unit/llms/codex tests/unit/llms/cohere tests/unit/llms/cometapi diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 1cea0194715..d4400fbb4cc 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -19,6 +19,7 @@ from litellm.a2a_protocol.utils import ( ) from litellm.exceptions import BadRequestError from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33 RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256 @@ -50,7 +51,7 @@ _RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = ( class _SSELineSource(Protocol): """Minimal streaming-response surface used to read SSE lines.""" - def aiter_lines(self) -> AsyncIterator[str]: ... + def aiter_text(self) -> AsyncIterator[str]: ... def _filter_reserved_headers( @@ -244,7 +245,7 @@ class BedrockAgentCoreA2ATransformation: Yields: Parsed JSON dicts from SSE data lines """ - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): line = line.strip() if not line: continue diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 0a2871d25cc..3937c47ea3c 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.custom_http import httpxSpecialProvider _IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token" @@ -74,7 +75,7 @@ class _WXORun(TypedDict, total=False): class _SSELineSource(Protocol): - def aiter_lines(self) -> AsyncIterator[str]: ... + def aiter_text(self) -> AsyncIterator[str]: ... class _WXOView(TypedDict, total=False): @@ -231,7 +232,7 @@ class WatsonxOrchestrateHandler: async def _accumulate_wxo_sse_text(response: _SSELineSource) -> str: source: Final[_WXOView] = {"sse_source": response} accumulated_text = "" - async for line in source["sse_source"].aiter_lines(): + async for line in aiter_sse_lines(source["sse_source"].aiter_text()): if not line.startswith("data:"): continue data_str = line[5:].strip() diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 47a3a903d82..ee3639ac387 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -34,6 +34,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool from litellm.responses.sse_output_recovery import ( parse_sse_json_chunk, @@ -1351,7 +1352,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): recovered_output_items: Final[dict[int, dict[str, object]]] = {} recovered_text_only_items: Final[dict[int, dict[str, object]]] = {} - for chunk in raw_sse.splitlines(): + for chunk in split_sse_lines(raw_sse): parsed_chunk = parse_sse_json_chunk(chunk) if parsed_chunk is None: continue diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index a22c47cb1e7..08ac7d4f9dc 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -131,9 +132,9 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps + # Gemini streamGenerateContent uses SSE line framing; iter_sse_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.iter_lines() + self.stream_iterator = iter_sse_lines(response.iter_text()) def __iter__(self): return self @@ -182,9 +183,9 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps + # Gemini streamGenerateContent uses SSE line framing; aiter_sse_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.aiter_lines() + self.stream_iterator = aiter_sse_lines(response.aiter_text()) def __aiter__(self): return self diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index f4ad57bf7b4..c7faa3fc510 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.interactions import ( InteractionsAPIStreamingResponse, ) @@ -145,7 +146,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator): litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) - self.stream_iterator = response.aiter_lines() + self.stream_iterator = aiter_sse_lines(response.aiter_text()) def __aiter__(self): return self @@ -212,7 +213,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) - self.stream_iterator = response.iter_lines() + self.stream_iterator = iter_sse_lines(response.iter_text()) def __iter__(self): return self diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3fb25de43cc..2348d8e57d6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1263,7 +1263,7 @@ class CustomStreamWrapper: ) if _is_delta_empty: - model_response.choices[0].delta = Delta(content=None) # ensure empty delta chunk returned + model_response.choices[0].delta = Delta(content=self._close_open_think_block()) # get any function call arguments model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason @@ -1292,21 +1292,26 @@ class CustomStreamWrapper: if self.merge_reasoning_content_in_choices is not True: return delta: Final = model_response.choices[0].delta - reasoning_content: Final = getattr(delta, "reasoning_content", None) + reasoning_content: Final = getattr(delta, "reasoning_content", None) or "" content: Final = delta.content or "" - if reasoning_content: - self._emitted_disqualifying_content = True - closes_thinking: Final = bool(content) and not self.sent_last_thinking_block - opening: Final = "" if self.sent_first_thinking_block else "" - delta.content = opening + reasoning_content + ("" + content if closes_thinking else "") - self.sent_first_thinking_block = True - self.sent_last_thinking_block = self.sent_last_thinking_block or closes_thinking - elif self.sent_first_thinking_block and not self.sent_last_thinking_block and content: - delta.content = "" + content - self.sent_last_thinking_block = True - + ends_thinking: Final = bool(content or delta.tool_calls or delta.function_call) + inside_think_block: Final = self.sent_first_thinking_block and not self.sent_last_thinking_block + opening: Final = "" if reasoning_content and not inside_think_block else "" + closing: Final = "" if ends_thinking and (reasoning_content or inside_think_block) else "" if hasattr(delta, "reasoning_content"): del delta.reasoning_content + if not reasoning_content and not closing: + return + delta.content = opening + reasoning_content + closing + content + self.sent_first_thinking_block = True + self.sent_last_thinking_block = bool(closing) + self._emitted_disqualifying_content = self._emitted_disqualifying_content or bool(reasoning_content) + + def _close_open_think_block(self) -> str | None: + if not self.sent_first_thinking_block or self.sent_last_thinking_block: + return None + self.sent_last_thinking_block = True + return "" def _dispatch_provider_chunk( self, @@ -1932,8 +1937,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason == "stop" and self.tool_call ): # don't overwrite for other - potential error finish reasons model_response.choices[0].finish_reason = "tool_calls" + close_tag: Final = self._close_open_think_block() + if close_tag is not None: + model_response.choices[0].delta.content = close_tag return model_response + def _final_chunk_for_logging(self, processed_chunk: "ModelResponseStream") -> "ModelResponseStream | None": + if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: + return processed_chunk + if processed_chunk.choices[0].delta.content is None: + return None + return processed_chunk.model_copy( + update={"choices": [processed_chunk.choices[0].model_copy(update={"finish_reason": None})]} + ) + def _record_usage_only_chunk(self, model_response: "ModelResponseStream") -> None: """ Keep provider usage-only chunks (e.g. OpenRouter's post-finish chunk, which carries a @@ -2173,10 +2190,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() - # The logged response is built from self.chunks; keep a finish_reason the provider sent on its - # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. - if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: - self.chunks.append(processed_chunk) + final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk) + if final_chunk_for_logging is not None: + self.chunks.append(final_chunk_for_logging) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) ## LOGGING @@ -2406,10 +2422,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() - # The logged response is built from self.chunks; keep a finish_reason the provider sent on its - # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. - if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: - self.chunks.append(processed_chunk) + final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk) + if final_chunk_for_logging is not None: + self.chunks.append(final_chunk_for_logging) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) # see sync __next__'s sibling branch: deliberately do NOT restore diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 13ceea3e31b..8a82ed98991 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -38,6 +38,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.llms.anthropic import ( ContentBlockDelta, ContentBlockStart, @@ -233,7 +234,7 @@ async def make_call( raise AnthropicError(status_code=500, message=str(e)) completion_stream: Final = ModelResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, json_mode=json_mode, speed=speed, @@ -304,7 +305,7 @@ def make_sync_call( ) completion_stream: Final = ModelResponseIterator( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, json_mode=json_mode, speed=speed, diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py index 90bdad8f4bc..c103eb0f975 100644 --- a/litellm/llms/anthropic/messages/stream_assembly.py +++ b/litellm/llms/anthropic/messages/stream_assembly.py @@ -13,6 +13,7 @@ from litellm.llms.anthropic.chat.handler import ( ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.messages.native.utils import MESSAGE_DELTA_OPTIONAL_KEYS +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.utils import ( Choices, Message, @@ -269,7 +270,7 @@ def build_complete_streaming_response_legacy( def extract_sse_data(event_str: str) -> dict | None: """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" - for line in event_str.splitlines(): + for line in split_sse_lines(event_str): stripped = line.strip() if stripped.startswith("data:"): payload = stripped[len("data:") :].strip() diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 94f1399f01b..423a09c5c92 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -35,6 +35,7 @@ from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionAnnotationURLCitation, @@ -632,7 +633,7 @@ class AzureAIAgentsHandler: current_event = None - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): line = line.strip() if line.startswith("event:"): diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 032951be74a..3afc685b3dd 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -96,7 +96,7 @@ class BaseModelResponseIterator: released (and a backend like vLLM aborts generation) when the stream is abandoned before its natural end. - ``streaming_response`` is usually a bare ``aiter_lines()`` generator + ``streaming_response`` is usually a bare ``aiter_sse_lines()`` generator that holds no reference to the response, so the handler that owns the response attaches it here after construction.""" if self.http_response is not None: diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 725146fe894..fb886e42780 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -6,6 +6,7 @@ from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row ha from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -24,7 +25,7 @@ def _anthropic_stream_chunk_events(item: object) -> list[dict]: events: Final[list[dict]] = [] for block in chunk.split("\n\n"): - for line in block.splitlines(): + for line in split_sse_lines(block): stripped = line.strip() if not stripped.startswith("data:"): continue diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index fe0e3abb87e..9b9b4e0a947 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -12,6 +12,7 @@ from typing import Final import httpx from pydantic import Field, PrivateAttr +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.base import LiteLLMPydanticObjectBase SANDBOX_MAX_OUTPUT_BYTES: Final = 10 * 1024 * 1024 @@ -99,7 +100,7 @@ class BaseSandboxConfig: async def _read_capped_lines(self, response: httpx.Response) -> list[str]: lines: Final[list[str]] = [] total = 0 - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): total += len(line.encode("utf-8")) if total > SANDBOX_MAX_OUTPUT_BYTES: raise ValueError( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 8c4702bcf99..0624c2c65db 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -49,6 +49,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsAuthParams, BedrockBatchRecordKind from litellm.types.llms.openai import ( @@ -400,7 +401,7 @@ def _batch_input_line(line_number: int, line: str) -> _BatchInputLine: def _non_blank_lines(file_content: str) -> Iterator[tuple[int, str]]: - numbered_lines: Final = enumerate(file_content.splitlines(), start=1) + numbered_lines: Final = enumerate(split_sse_lines(file_content), start=1) return ((line_number, line) for line_number, line in numbered_lines if line.strip()) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 1b01b0ec6d0..3a65a8928bf 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -45,6 +45,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits imp structured_edits, tool_edits, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.guardrails.auto_router_compression import is_compression_guardrail from litellm.types.utils import GenericGuardrailAPIInputs @@ -947,7 +948,7 @@ def _invoke_frames(items: Sequence[object]) -> tuple[InvokeStreamFrame, ...]: def _sse_events(sse: bytes) -> tuple[Mapping[str, object], ...]: return tuple( _JSON_OBJECT.validate_python(json.loads(line[len("data:") :])) - for line in sse.decode().splitlines() + for line in split_sse_lines(sse.decode()) if line.startswith("data:") ) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 565fc12aa1e..6f72cadad67 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -47,6 +47,7 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.secret_managers.main import get_secret_str # AgentCore web-search rejects queries longer than 200 characters @@ -146,7 +147,9 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: progress notifications before the JSON-RPC response. """ for chunk in _SSE_EVENT_SEPARATOR.split(text): - payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:")) + payload = "\n".join( + line[len("data:") :].lstrip() for line in split_sse_lines(chunk) if line.startswith("data:") + ) if not payload: continue try: diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 040b7b8c1d8..72b5cde21d1 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo safe_convert_created_field, ) from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.sse_output_recovery import ( @@ -199,7 +200,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): error_message = None streamed_output_items: Final[dict[int, dict[str, object]]] = {} text_only_output_items: Final[dict[int, dict[str, object]]] = {} - for chunk in body_text.splitlines(): + for chunk in split_sse_lines(body_text): parsed_chunk = parse_sse_json_chunk(chunk) if parsed_chunk is None: continue diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 6bebfe2d444..6b172e77295 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.utils import TextChoices from litellm.utils import CustomStreamWrapper, TextCompletionResponse @@ -109,7 +110,7 @@ async def make_call( if response.status_code != 200: raise TextCompletionCodestralError(status_code=response.status_code, message=response.text) - completion_stream: Final = response.aiter_lines() + completion_stream: Final = aiter_sse_lines(response.aiter_text()) # LOGGING logging_obj.post_call( input=messages, @@ -339,13 +340,13 @@ class CodestralTextCompletion: response = litellm.module_level_client.post( completion_url, headers=headers, - data=json.dumps(data), + data=json.dumps({**data, "stream": True}), stream=stream, ) _response: Final = CustomStreamWrapper( - response.iter_lines(), + iter_sse_lines(response.iter_text()), model, - custom_llm_provider="codestral", + custom_llm_provider="text-completion-codestral", logging_obj=logging_obj, ) return _response diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 447d7b685f7..4060d84f7f4 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, get_ssl_configuration, ) +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.types.llms.openai import FileTypes from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProviders from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager @@ -466,7 +467,7 @@ class BaseLLMAIOHTTPHandler: ) else: completion_stream = provider_config.get_model_response_iterator( - streaming_response=response.iter_lines(), sync_stream=True + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True ) # LOGGING diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index afa12b58b29..e7da6586b97 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -118,6 +118,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, header_value, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, @@ -1086,7 +1087,7 @@ class BaseLLMHTTPHandler: completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator_for_request( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), @@ -1229,7 +1230,7 @@ class BaseLLMHTTPHandler: completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator_for_request( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), diff --git a/litellm/llms/custom_httpx/sse_lines.py b/litellm/llms/custom_httpx/sse_lines.py new file mode 100644 index 00000000000..71ea7d621eb --- /dev/null +++ b/litellm/llms/custom_httpx/sse_lines.py @@ -0,0 +1,74 @@ +""" +Split a decoded Server-Sent Events body into lines the way the SSE spec does. + +httpx's `iter_lines()` / `aiter_lines()` follow `str.splitlines()`, which also breaks on U+0085, U+2028, +U+2029, \\x0b, \\x0c and \\x1c-\\x1e. JSON leaves U+0085, U+2028 and U+2029 unescaped inside strings, so a +`data:` line carrying one of them gets cut in half. SSE lines end only at CRLF, LF or CR +(https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream) +""" + +import io +import re +from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence +from typing import Final + +_SSE_LINE_BREAK: Final = re.compile(r"\r\n|\r|\n") +_SSE_LINE_WITH_BREAK: Final = re.compile(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+") + + +def _split_lines(text: str) -> Sequence[str]: + if "\r" not in text: + return text.split("\n") + crlf_folded: Final = text.replace("\r\n", "\n") + if "\r" not in crlf_folded: + return crlf_folded.split("\n") + return _SSE_LINE_BREAK.split(text) + + +class _SSELineSplitter: + def __init__(self) -> None: + self._partial_line: io.StringIO = io.StringIO() + self._after_cr: bool = False + + def feed(self, text: str) -> tuple[str, ...]: + if not text: + return () + fresh: Final = text[1:] if self._after_cr and text.startswith("\n") else text + self._after_cr = fresh.endswith("\r") + *ended, tail = _split_lines(fresh) + if not ended: + _ = self._partial_line.write(tail) + return () + first: Final = self._partial_line.getvalue() + ended[0] + self._partial_line = io.StringIO() + _ = self._partial_line.write(tail) + return (first, *ended[1:]) + + def flush(self) -> tuple[str, ...]: + rest: Final = self._partial_line.getvalue() + self._partial_line = io.StringIO() + return (rest,) if rest else () + + +def iter_sse_lines(text_chunks: Iterable[str]) -> Iterator[str]: + splitter: Final = _SSELineSplitter() + for text in text_chunks: + yield from splitter.feed(text) + yield from splitter.flush() + + +def split_sse_lines(text: str) -> tuple[str, ...]: + return tuple(iter_sse_lines((text,))) + + +def split_sse_lines_keepends(text: str) -> tuple[str, ...]: + return tuple(_SSE_LINE_WITH_BREAK.findall(text)) + + +async def aiter_sse_lines(text_chunks: AsyncIterable[str]) -> AsyncIterator[str]: + splitter: Final = _SSELineSplitter() + async for text in text_chunks: + for line in splitter.feed(text): + yield line + for line in splitter.flush(): + yield line diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py index 34007bc8fd4..b6b2aa6c28f 100644 --- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -23,6 +23,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamTransformSink, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -121,7 +122,7 @@ def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: return tuple( text - for line in sse_text.splitlines() + for line in split_sse_lines(sse_text) if line.startswith("data:") for payload in (line[len("data:") :].strip(),) if payload and payload != "[DONE]" diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 1b7d6483739..6e6a117227e 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -357,7 +358,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): api_key: str | None = None, ) -> list[dict[str, str]]: streamed_response: Final = CustomStreamWrapper( - completion_stream=response.iter_lines(), + completion_stream=iter_sse_lines(response.iter_text()), model=model, custom_llm_provider="huggingface", logging_obj=logging_obj, diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 575820bceee..6fe355ad98c 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -12,6 +12,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices _DECODED_JSON: Final = TypeAdapter(object) @@ -36,12 +37,12 @@ class LangGraphSSEStreamIterator: def __iter__(self): """Initialize sync iteration.""" - self.line_iterator = self.response.iter_lines() + self.line_iterator = iter_sse_lines(self.response.iter_text()) return self def __aiter__(self): """Initialize async iteration.""" - self.async_line_iterator = self.response.aiter_lines() + self.async_line_iterator = aiter_sse_lines(self.response.aiter_text()) return self def _parse_sse_line(self, line: str) -> ModelResponseStream | None: diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 8879c20c5eb..b4dc57ff5ed 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final import litellm from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_httpx_client +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.utils import EmbeddingResponse, ModelResponse, Usage from ..common_utils import OobaboogaError @@ -74,7 +75,7 @@ def completion( stream=optional_params["stream"] if "stream" in optional_params else False, ) if "stream" in optional_params and optional_params["stream"] is True: - return response.iter_lines() + return iter_sse_lines(response.iter_text()) else: return oobabooga_config.transform_response( model=model, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 2c74ef19e7d..a450947afdf 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -16,6 +16,7 @@ from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.databricks.streaming_utils import ModelResponseIterator from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai.openai import OpenAIConfig @@ -65,7 +66,9 @@ async def make_call( model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False) + completion_stream = ModelResponseIterator( + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False + ) # LOGGING logging_obj.post_call( input=messages, @@ -103,7 +106,9 @@ def make_sync_call( model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True) + completion_stream = ModelResponseIterator( + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True + ) # LOGGING logging_obj.post_call( diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py index d0a61ea6e00..2ba4a63e08b 100644 --- a/litellm/llms/predibase/chat/handler.py +++ b/litellm/llms/predibase/chat/handler.py @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.utils import CustomStreamWrapper, ModelResponse from ..common_utils import PredibaseError @@ -34,7 +35,7 @@ async def make_call( if response.status_code != 200: raise PredibaseError(status_code=response.status_code, message=response.text) - completion_stream: Final = response.aiter_lines() + completion_stream: Final = aiter_sse_lines(response.aiter_text()) # LOGGING logging_obj.post_call( input=messages, @@ -163,7 +164,7 @@ class PredibaseChatCompletion: timeout=timeout, ) _response: Final = CustomStreamWrapper( - response.iter_lines(), + iter_sse_lines(response.iter_text()), model, custom_llm_provider="predibase", logging_obj=logging_obj, diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 9f592398f74..0339533e1f5 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) @@ -357,7 +358,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) -> VertexAgentEngineResponseIterator: """Return a streaming iterator for SSE responses.""" return VertexAgentEngineResponseIterator( - streaming_response=raw_response.iter_lines(), + streaming_response=iter_sse_lines(raw_response.iter_text()), sync_stream=True, ) @@ -465,7 +466,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): # Create iterator for SSE stream (async) completion_stream: Final = VertexAgentEngineResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index f27593dc94e..c506192f253 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -41,6 +41,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.llms.anthropic import AnthropicThinkingParam from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.gemini import BidiGenerateContentServerMessage @@ -3241,7 +3242,7 @@ async def make_call( raise completion_stream: Final = ModelResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, logging_obj=logging_obj, response_headers=response.headers, @@ -3312,7 +3313,7 @@ def make_sync_call( raise completion_stream: Final = ModelResponseIterator( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, logging_obj=logging_obj, response_headers=response.headers, @@ -4190,7 +4191,7 @@ class ModelResponseIterator: log_failed_hop(e) return self._truncated_stream_end(token) self._start_hop(hop_response, token) - self.streaming_response = hop_response.aiter_lines() + self.streaming_response = aiter_sse_lines(hop_response.aiter_text()) self.async_response_iterator = self.streaming_response return None @@ -4204,7 +4205,7 @@ class ModelResponseIterator: log_failed_hop(e) return self._truncated_stream_end(token) self._start_hop(hop_response, token) - self.streaming_response = hop_response.iter_lines() + self.streaming_response = iter_sse_lines(hop_response.iter_text()) self.response_iterator = self.streaming_response return None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 4691c7b0bc5..219ce906340 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( A2AVersion, @@ -320,7 +321,7 @@ async def _a2a_sse_event_source( "error": {"code": -32603, "message": resp.reason_phrase}, } return - async for line in resp.aiter_lines(): + async for line in aiter_sse_lines(resp.aiter_text()): stripped = line.strip() if not stripped.startswith("data:"): continue diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index f748a754fe0..32e0f6e8682 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -14,6 +14,8 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines_keepends + _MESSAGE_START_EVENT: Final = "message_start" _MESSAGE_START_MARKER: Final = b"message_start" _SSE_DATA_FIELD: Final = "data:" @@ -52,7 +54,7 @@ def _restamped_data_line(line: str, requested_model: str) -> str | None: def _restamped_frame(frame: str, requested_model: str) -> str | None: - lines: Final = frame.splitlines(keepends=True) + lines: Final = split_sse_lines_keepends(frame) restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) if all(line is None for line in restamped): return None diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 7c7eb51bc5d..5415314f29b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -29,6 +29,7 @@ from litellm.llms.anthropic.messages.stream_assembly import ( split_sse_chunk_into_events, ) from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -61,7 +62,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i def _event_payload(event: str) -> Mapping[str, object] | None: data_lines: Final = tuple( - line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + line.strip()[len("data:") :].strip() for line in split_sse_lines(event) if line.strip().startswith("data:") ) if not data_lines or data_lines[0] in ("", "[DONE]"): return None @@ -362,7 +363,7 @@ class AnthropicPassthroughLoggingHandler: ) -> str | None: for raw in all_chunks: text = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for line in text.splitlines(): + for line in split_sse_lines(text): if not line.startswith("data:"): continue try: @@ -402,7 +403,7 @@ class AnthropicPassthroughLoggingHandler: """ for raw in reversed(all_chunks): text = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for line in reversed(text.splitlines()): + for line in reversed(split_sse_lines(text)): if not line.startswith("data:"): continue try: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py index 2afc1d17bc2..0c4a3f64852 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.tinyfish import ( @@ -81,7 +82,7 @@ def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool: def run_id_from_sse_frames(frames: bytes) -> str | None: - return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines()) + return _run_id_from_sse_chunks(split_sse_lines(frames.decode("utf-8", errors="replace"))) def resolve_tinyfish_agent_api_base() -> str: diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 9790c150ae1..cae8ccc93ee 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -51,6 +51,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( can_access_resource, resolve_resource_owner_id, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames @@ -830,7 +831,7 @@ _SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) def _first_streamed_response(frames: bytes) -> tuple[str, Mapping[str, JsonValue]] | None: - for line in frames.decode("utf-8", errors="replace").splitlines(): + for line in split_sse_lines(frames.decode("utf-8", errors="replace")): if not line.startswith(_SSE_DATA_PREFIX): continue try: diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 5f6a7432afb..1eb76f8f558 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -18,6 +18,7 @@ from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -81,7 +82,7 @@ class _StreamsBody(Protocol): def _sse_frame_data(frame: str) -> str | None: - return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) + return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None) async def _never_receive() -> Message: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 219944de521..982158040c4 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -165,6 +165,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._upstream_exhausted: bool = False self._response_id_primed: bool = False self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] + self._tool_calls_held_behind_text: tuple[object, ...] = () self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream @@ -388,6 +389,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return "" return self._tool_call_id_by_index.get(index, "") + def _hold_tool_calls_behind_text(self, tool_calls: object) -> None: + if isinstance(tool_calls, list): + self._tool_calls_held_behind_text = (*self._tool_calls_held_behind_text, *tool_calls) + + def _release_held_tool_calls(self) -> tuple[object, ...]: + held: Final = self._tool_calls_held_behind_text + self._tool_calls_held_behind_text = () + return held + def _queue_tool_call_delta_events(self, tool_calls: object) -> None: """ Convert chat-completions streaming `tool_calls` deltas into Responses API streaming events. @@ -1365,6 +1375,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if delta_content: if not self.sent_message_item_added_event: self._queue_message_item_added_events() + self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None)) self._sequence_number += 1 text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, @@ -1379,7 +1390,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 3: Handle tool call deltas (if any) -> queue events and emit them # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior if chunk.choices and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls: - self._queue_tool_call_delta_events(chunk.choices[0].delta.tool_calls) + self._queue_tool_call_delta_events([*self._release_held_tool_calls(), *chunk.choices[0].delta.tool_calls]) # Return one pending tool event at a time if self._pending_tool_events: return self._pending_tool_events.pop(0) diff --git a/tests/benchmarks/test_sse_lines_benchmarks.py b/tests/benchmarks/test_sse_lines_benchmarks.py new file mode 100644 index 00000000000..1320fae49a6 --- /dev/null +++ b/tests/benchmarks/test_sse_lines_benchmarks.py @@ -0,0 +1,27 @@ +from collections.abc import Sequence +from typing import Final + +import pytest + +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines + +EVENT: Final = 'data: {"choices":[{"delta":{"content":"hello world"}}]}' + + +def _chunks(line_break: str) -> Sequence[str]: + body: Final = "".join(f"{EVENT}{line_break}{line_break}" for _ in range(2000)) + return tuple(body[i : i + 512] for i in range(0, len(body), 512)) + + +LF_CHUNKS: Final = _chunks("\n") +CRLF_CHUNKS: Final = _chunks("\r\n") + + +@pytest.mark.benchmark +def test_iter_sse_lines_lf_stream() -> None: + assert sum(1 for _ in iter_sse_lines(LF_CHUNKS)) == 4000 + + +@pytest.mark.benchmark +def test_iter_sse_lines_crlf_stream() -> None: + assert sum(1 for _ in iter_sse_lines(CRLF_CHUNKS)) == 4000 diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 09f09170241..cd3f7622fa2 100644 --- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -10,6 +10,7 @@ Verifies that: """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -560,6 +561,37 @@ class TestStreaming: assert sent_body["params"]["message"]["messageId"] == "msg-001" assert [event["result"]["kind"] for event in events] == ["task", "status-update"] + @pytest.mark.asyncio + @pytest.mark.usefixtures("httpx_transport") + async def test_streaming_keeps_unicode_line_separators_inside_an_sse_event(self) -> None: + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + event: Final = { + "jsonrpc": "2.0", + "id": "req-001", + "result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "a\u2028b\u2029c"}]}}, + } + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(event, ensure_ascii=False)}\n\n", + ) + ) + events: Final = [ + streamed + async for streamed in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=SAMPLE_PARAMS, + litellm_params=SAMPLE_LITELLM_PARAMS, + ) + ] + + assert events == [event] + class TestConfigManager: """Test that config manager routes 'bedrock' correctly.""" diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index a300560ae9d..ac3f0390f84 100644 --- a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,6 +1,7 @@ import asyncio import json import time +from typing import Final import httpx import pytest @@ -36,13 +37,8 @@ class _ShortTtlTokenClient: return _JsonResponse({"access_token": f"token-{self.calls}", "expires_in": 30}) -class _SSELines: - def __init__(self, lines): - self.lines = lines - - async def aiter_lines(self): - for line in self.lines: - yield line +def _sse_response(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) class _InvalidJsonStreamResponse: @@ -218,7 +214,7 @@ def test_cp4d_token_ttl_from_absolute_expiration(): @pytest.mark.asyncio async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): - response = _SSELines( + response = _sse_response( [ "data: null", "data: true", @@ -230,6 +226,14 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): ) +@pytest.mark.asyncio +async def test_accumulate_wxo_sse_text_keeps_unicode_line_separators_inside_an_event() -> None: + separated_text: Final = "before\u2028between\u2029after" + response: Final = _sse_response([f"data: {json.dumps({'results': separated_text}, ensure_ascii=False)}", ""]) + + assert await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) == separated_text + + @pytest.mark.asyncio async def test_short_lived_tokens_are_not_served_from_cache(): client = _ShortTtlTokenClient() diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 9aa494c1aab..46f1dbf2e7a 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1102,6 +1102,24 @@ def test_recover_output_items_merges_text_only_items_at_distinct_indices(): assert recovered[1]["content"][0]["text"] == "From OUTPUT_TEXT_DONE only" +def test_recover_output_items_keeps_an_item_whose_text_holds_unicode_line_separators() -> None: + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + item: Final = { + "type": "message", + "id": "msg_item_0", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "From\u2028OUTPUT_ITEM\u2029DONE", "annotations": []}], + } + event: Final = {"type": "response.output_item.done", "output_index": 0, "item": item} + raw_sse: Final = f"data: {json.dumps(event, ensure_ascii=False)}\ndata: [DONE]\n" + + assert LiteLLMResponsesTransformationHandler._recover_output_items_from_raw_sse(raw_sse) == [item] + + def test_transform_response_prefers_completed_output_from_raw_sse(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, diff --git a/tests/unit/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py index 120c99b6e42..cbb9638b3d9 100644 --- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py @@ -62,7 +62,9 @@ def _logging_obj( ], ) def test_streaming_logging_targets_the_provider_that_served_the_request( - iterator_cls: type, + iterator_cls: ( + type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator] + ), custom_llm_provider: str, expected_endpoint_type: EndpointType, ): @@ -99,18 +101,15 @@ def _large_inline_data_event() -> str: return f"data: {json.dumps(payload)}" +def _sse_response(body: str) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream(body.encode())) + + @pytest.mark.asyncio async def test_async_streaming_iterator_yields_complete_sse_events(): """Large inlineData must not be split across byte-chunk boundaries.""" - mock_response = MagicMock() - - async def _aiter_lines(): - yield _large_inline_data_event() - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(_large_inline_data_event() + "\n\n"), model="gemini-3.1-flash-image-preview", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -128,11 +127,8 @@ async def test_async_streaming_iterator_yields_complete_sse_events(): def test_sync_streaming_iterator_yields_complete_sse_events(): - mock_response = MagicMock() - mock_response.iter_lines.return_value = iter([_large_inline_data_event()]) - iterator = GoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(_large_inline_data_event() + "\n\n"), model="gemini-3.1-flash-image-preview", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -150,17 +146,8 @@ def test_sync_streaming_iterator_yields_complete_sse_events(): @pytest.mark.asyncio async def test_async_streaming_iterator_preserves_multi_field_sse_event(): - mock_response = MagicMock() - - async def _aiter_lines(): - yield "event: message" - yield 'data: {"text":"hi"}' - yield "" - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response('event: message\ndata: {"text":"hi"}\n\n'), model="gemini-test", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -174,16 +161,8 @@ async def test_async_streaming_iterator_preserves_multi_field_sse_event(): @pytest.mark.asyncio async def test_async_streaming_iterator_forwards_sse_comment_events(): - mock_response = MagicMock() - - async def _aiter_lines(): - yield ": keepalive" - yield "" - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(": keepalive\n\n"), model="gemini-test", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -195,6 +174,46 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): assert chunk == b": keepalive\n\n" +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _first_chunk( + iterator: GoogleGenAIGenerateContentStreamingIterator | AsyncGoogleGenAIGenerateContentStreamingIterator, +) -> bytes: + if isinstance(iterator, AsyncGoogleGenAIGenerateContentStreamingIterator): + return await iterator.__anext__() + return next(iterator) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "iterator_cls", + [AsyncGoogleGenAIGenerateContentStreamingIterator, GoogleGenAIGenerateContentStreamingIterator], +) +async def test_streaming_iterator_keeps_unicode_line_separators_inside_the_event( + iterator_cls: ( + type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator] + ), +) -> None: + payload: Final = {"candidates": [{"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}]}}]} + event: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + iterator: Final = iterator_cls( + response=_sse_response(event), + model="gemini-3.8-flash", + logging_obj=_logging_obj(), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk: Final = await _first_chunk(iterator) + + assert chunk == event.encode() + assert ( + json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0]["text"] == _LINE_SEPARATED_TEXT + ) + + def _gemini_event(text: str, finish_reason: str | None = None) -> str: candidate: Final = {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0} payload: Final = { diff --git a/tests/unit/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py index afcd0d559d7..fb1fe92d637 100644 --- a/tests/unit/interactions/test_interactions_streaming_iterator.py +++ b/tests/unit/interactions/test_interactions_streaming_iterator.py @@ -16,7 +16,10 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.interactions import streaming_iterator as interactions_streaming_iterator_module -from litellm.interactions.streaming_iterator import InteractionsAPIStreamingIterator +from litellm.interactions.streaming_iterator import ( + InteractionsAPIStreamingIterator, + SyncInteractionsAPIStreamingIterator, +) from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig @@ -148,3 +151,43 @@ async def test_stream_logs_its_first_event_time_measured_from_the_call_start(): assert first_event_at is not None assert recorder.logged == [(call_start, first_event_at)] + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _first_event( + iterator: InteractionsAPIStreamingIterator | SyncInteractionsAPIStreamingIterator, +) -> InteractionsAPIStreamingResponse: + if isinstance(iterator, InteractionsAPIStreamingIterator): + return await iterator.__anext__() + return next(iterator) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterator_cls", [InteractionsAPIStreamingIterator, SyncInteractionsAPIStreamingIterator]) +async def test_stream_keeps_unicode_line_separators_inside_the_delta_text( + iterator_cls: type[InteractionsAPIStreamingIterator] | type[SyncInteractionsAPIStreamingIterator], +) -> None: + delta: Final = {"type": "text", "text": _LINE_SEPARATED_TEXT} + event: Final = {"event_type": "step.delta", "id": "interaction-1", "delta": delta} + iterator: Final = iterator_cls( + response=httpx.Response( + 200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode()) + ), + model="gemini/gemini-3.1-pro-preview", + interactions_api_config=GoogleAIStudioInteractionsConfig(), + logging_obj=LitellmLogging( + model="gemini/gemini-3.1-pro-preview", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="ainteraction", + start_time=datetime(2026, 10, 9, 12, 0, 0), + litellm_call_id="interactions-separators", + function_id="interactions-separators", + ), + ) + + first: Final = await _first_event(iterator) + + assert first.delta == delta diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index a1b58521d7e..79978ecfe8e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -506,6 +506,118 @@ def test_multi_chunk_reasoning_and_content( assert initialized_custom_stream_wrapper.sent_last_thinking_block is True +def _stream_chunk(delta: dict[str, object], finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason}]) + + +def _content_through_first_call(deltas: Sequence[Delta]) -> str: + first_call: Final = next( + (index for index, delta in enumerate(deltas) if delta.tool_calls or delta.function_call), len(deltas) + ) + return "".join(delta.content or "" for delta in deltas[: first_call + 1]) + + +_GET_WEATHER_CALL: Final = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize( + "chunks,expected_content,expected_tool_calls", + [ + ( + ( + _stream_chunk({"reasoning_content": "Check the weather"}), + _stream_chunk({"tool_calls": [_GET_WEATHER_CALL]}), + _stream_chunk({}, "tool_calls"), + ), + "Check the weather", + 1, + ), + ( + ( + _stream_chunk({"reasoning_content": "Pick a function"}), + _stream_chunk({"function_call": {"name": "get_weather", "arguments": "{}"}}), + _stream_chunk({}, "function_call"), + ), + "Pick a function", + 0, + ), + ( + ( + _stream_chunk({"reasoning_content": "Step one"}), + _stream_chunk({"reasoning_content": ", step two"}), + _stream_chunk({}, "length"), + ), + "Step one, step two", + 0, + ), + ( + ( + _stream_chunk({"reasoning_content": "a"}), + _stream_chunk({"content": "b"}), + _stream_chunk({"reasoning_content": "c"}), + _stream_chunk({"content": "d"}), + _stream_chunk({}, "stop"), + ), + "abcd", + 0, + ), + ( + (_stream_chunk({"reasoning_content": "Still thinking"}),), + "Still thinking", + 0, + ), + ( + (_stream_chunk({"content": "hi"}), _stream_chunk({}, "stop")), + "hi", + 0, + ), + ], + ids=[ + "reasoning_then_tool_call", + "reasoning_then_function_call", + "reasoning_cut_by_length", + "reasoning_after_answer", + "stream_ends_mid_reasoning", + "answer_without_reasoning", + ], +) +async def test_merged_reasoning_stream_keeps_think_tags_balanced( + logging_obj: Logging, + is_async: bool, + chunks: tuple[ModelResponseStream, ...], + expected_content: str, + expected_tool_calls: int, +) -> None: + response: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=[chunk.model_copy(deep=True) for chunk in chunks]), + model="deepseek-reasoner", + logging_obj=logging_obj, + ) + response.merge_reasoning_content_in_choices = True + + deltas: Final = ( + tuple([chunk.choices[0].delta async for chunk in response]) + if is_async + else tuple(chunk.choices[0].delta for chunk in response) + ) + logged: Final = litellm.stream_chunk_builder(chunks=response.chunks) + + assert "".join(delta.content or "" for delta in deltas) == expected_content + assert _content_through_first_call(deltas) == expected_content + assert sum(len(delta.tool_calls or ()) for delta in deltas) == expected_tool_calls + assert logged.choices[0].message.content == (expected_content or None) + assert any(chunk.choices[0].finish_reason for chunk in response.chunks if chunk.choices) == any( + chunk.choices[0].finish_reason for chunk in chunks + ) + + def test_strip_sse_data_from_chunk(): """Test the static method that strips 'data: ' prefix from SSE chunks""" # Test with string inputs @@ -7849,7 +7961,7 @@ def gemini_mock_post_streaming(url, **kwargs): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "text/event-stream"} - mock_response.iter_lines = MagicMock(return_value=stream_response()) + mock_response.iter_text = MagicMock(return_value=stream_response()) return mock_response diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 1f4a5aaa256..355d407001c 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -15,7 +15,7 @@ import litellm from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call +from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( @@ -27,24 +27,34 @@ from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE +def _text_delta_stream(text: str) -> bytes: + events: Final = ( + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + return "".join( + f"event: {name}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n" for name, payload in events + ).encode() + + +def _streamed_text(chunks: list[object]) -> str: + return "".join(chunk.choices[0].delta.content or "" for chunk in chunks if isinstance(chunk, ModelResponseStream)) + + @pytest.mark.asyncio async def test_make_call_passes_logging_obj_to_client_post(): """make_call must pass logging_obj to client.post so track_llm_api_timing can set llm_api_duration_ms for litellm_overhead_time_ms.""" mock_client = AsyncMock() - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock( - return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']) - ) - mock_client.post.return_value = mock_response + mock_client.post.return_value = httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream("Hi there"))) logging_obj = MagicMock() - await make_call( + completion_stream, _ = await make_call( client=mock_client, api_base="https://api.anthropic.com/v1/messages", headers={}, data="{}", - model="claude-3-5-haiku", + model="claude-haiku-5-5", messages=[{"role": "user", "content": "Hi"}], logging_obj=logging_obj, timeout=60.0, @@ -54,6 +64,39 @@ async def test_make_call_passes_logging_obj_to_client_post(): mock_client.post.assert_called_once() call_kwargs = mock_client.post.call_args[1] assert call_kwargs.get("logging_obj") is logging_obj + assert _streamed_text([chunk async for chunk in completion_stream]) == "Hi there" + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _streamed_chunks(sync_stream: bool, text: str) -> list[object]: + transport: Final = httpx.MockTransport( + lambda request: httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream(text))) + ) + call_args: Final = { + "api_base": "https://api.anthropic.com/v1/messages", + "headers": {}, + "data": "{}", + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "Hi"}], + "logging_obj": MagicMock(), + "timeout": 60.0, + "json_mode": False, + } + if sync_stream: + sync_iterator, _ = make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args) + return list(sync_iterator) + async_iterator, _ = await make_call(client=AsyncHTTPHandler(transport=transport), **call_args) + return [chunk async for chunk in async_iterator] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_streamed_text_keeps_unicode_line_separators(sync_stream: bool) -> None: + chunks: Final = await _streamed_chunks(sync_stream, _LINE_SEPARATED_TEXT) + + assert _streamed_text(chunks) == _LINE_SEPARATED_TEXT def test_anthropic_completion_does_not_send_deployment_default_limits(): diff --git a/tests/unit/llms/anthropic/messages/test_stream_assembly.py b/tests/unit/llms/anthropic/messages/test_stream_assembly.py index f7d6b8d4f8c..4993d2fa9bf 100644 --- a/tests/unit/llms/anthropic/messages/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/messages/test_stream_assembly.py @@ -144,6 +144,30 @@ def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops( assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {} +def test_anthropic_message_delta_fields_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + safeguard_results: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes\u2028the\u2029root"} + }, + }, + } + ] + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 2}, + } + raw_message_delta_frame: Final = f"event: message_delta\ndata: {json.dumps(message_delta, ensure_ascii=False)}\n\n" + + fields: Final = anthropic_message_delta_fields((*ANTHROPIC_STREAM[:5], raw_message_delta_frame.encode())) + + assert dict(fields) == {"safeguard_results": safeguard_results} + + @pytest.mark.parametrize( "chunks", [GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()], diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py index 71466bb172c..14b6094747f 100644 --- a/tests/unit/llms/azure_ai/agents/test_handler.py +++ b/tests/unit/llms/azure_ai/agents/test_handler.py @@ -173,7 +173,7 @@ async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None: def _sse(event: str, data: dict[str, object] | str) -> str: - return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data)}\n\n" + return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)}\n\n" def _message_delta(text: str) -> dict[str, object]: @@ -222,3 +222,24 @@ async def test_acompletion_stream_yields_agent_message_deltas_then_stop(respx_mo "thread": {"messages": [{"role": "user", "content": "Hi Agent, what is 10 + 5?"}]}, } assert chunks == [("10 + 5", None), (" = 15", None), (None, "stop")] + + +@pytest.mark.asyncio +async def test_acompletion_stream_keeps_unicode_line_separators_inside_a_delta(respx_mock: respx.MockRouter) -> None: + separated_text: Final = "before\u2028between\u2029after" + _ = respx_mock.post(f"{API_BASE}/threads/runs").respond( + 200, + text=_sse("thread.message.delta", _message_delta(separated_text)) + _sse("done", "[DONE]"), + headers={"content-type": "text/event-stream"}, + ) + + response: Final = await litellm.acompletion( + model=f"azure_ai/agents/{AGENT_ID}", + messages=[{"role": "user", "content": "Say it with separators"}], + api_base=API_BASE, + api_key="azure-ad-token", + stream=True, + ) + chunks: Final = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response] + + assert chunks == [(separated_text, None), (None, "stop")] diff --git a/tests/unit/llms/azure_ai/agents/test_transformation.py b/tests/unit/llms/azure_ai/agents/test_transformation.py index 10676438858..413908c8313 100644 --- a/tests/unit/llms/azure_ai/agents/test_transformation.py +++ b/tests/unit/llms/azure_ai/agents/test_transformation.py @@ -1,7 +1,7 @@ import json from typing import Final -from unittest.mock import MagicMock +import httpx import pytest from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig @@ -428,15 +428,10 @@ async def test_azure_ai_agents_streaming_accumulates_annotations_from_multiple_t "data: [DONE]", ] - async def mock_aiter_lines(): - for line in sse_lines: - yield line - - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines()) + sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode())) chunks = [] - async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"): + async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"): chunks.append(chunk) final_chunk = chunks[-1] @@ -498,15 +493,10 @@ async def test_azure_ai_agents_streaming_annotations_from_completed_message(): "data: [DONE]", ] - async def mock_aiter_lines(): - for line in sse_lines: - yield line - - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines()) + sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode())) chunks = [] - async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"): + async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"): chunks.append(chunk) assert len(chunks) >= 1 diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py index 65208a2bd3b..acd50616f90 100644 --- a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py +++ b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py @@ -1,9 +1,11 @@ +import json from typing import Final import pytest import litellm from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, ) @@ -32,3 +34,18 @@ def test_a_per_guardrail_skip_flag_wins_over_the_global_setting( assert effective_skip_system_message_for_guardrail(guardrail) is expected assert effective_skip_tool_message_for_guardrail(guardrail) is expected + + +def test_blocked_response_usage_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + message_start: Final = {"type": "message_start", "message": {"usage": {"input_tokens": 11, "output_tokens": 1}}} + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"}, + "usage": {"output_tokens": 7}, + } + stream: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n" + for event in (message_start, message_delta) + ).encode() + + assert blocked_response_usage([stream]) == {"input_tokens": 11, "output_tokens": 7} diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 7e05ef500f5..9d9e39df944 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -4458,6 +4458,20 @@ class TestMalformedBatchJsonl: _VALID_BATCH_RECORD: Final = json.loads(_VALID_BATCH_LINE) +def test_file_content_conversion_keeps_a_line_whose_json_holds_unicode_line_separators() -> None: + text: Final = "Hello\u2028from\u2029batch" + body: Final = {**_VALID_BATCH_RECORD["body"], "messages": [{"role": "user", "content": text}]} + record: Final = {**_VALID_BATCH_RECORD, "body": body} + file_content: Final = f"{json.dumps(record, ensure_ascii=False)}\n".encode() + + bedrock_jsonl, _ = BedrockJsonlFilesTransformation().transform_openai_file_content_to_bedrock_file_content( + file_content + ) + + (bedrock_record,) = (json.loads(line) for line in bedrock_jsonl.split("\n")) + assert bedrock_record["modelInput"]["messages"] == [{"role": "user", "content": [{"text": text}]}] + + def _batch_line(body: object) -> str: return json.dumps({**_VALID_BATCH_RECORD, "body": body}) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 5aaee1f0ee4..aec1672b335 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -18,7 +18,7 @@ import types from collections.abc import Callable, Mapping, Sequence from datetime import datetime from itertools import chain -from typing import ClassVar, Literal, TypeAlias +from typing import ClassVar, Final, Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -60,6 +60,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.handler import ( BedrockPassthroughGuardrailHandler, _extract_converse_texts, _is_converse_endpoint, + _sse_events, _write_back_texts, ) @@ -3840,3 +3841,14 @@ class TestInvokeStreamGuardrailRelay: assert received_key is not None and received_key.texts == ("In Paris",) assert rewritten_key == received_key + + +def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None: + event: Final = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Card\u2028on\u2029file"}, + } + sse: Final = f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + assert _sse_events(sse) == (event,) diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py index 664dbf87fe9..01cca1fc270 100644 --- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py @@ -8,6 +8,7 @@ the sharded CI (coverage collection runs against this tree). import json import os +from typing import Final import httpx import pytest @@ -682,3 +683,17 @@ def test_transform_search_response_ignores_text_blocks_without_result_objects(bl body = {"result": {"content": [{"type": "text", "text": block_text}]}} assert _transform_text(json.dumps(body)).results == [] + + +def test_transform_search_response_keeps_a_result_whose_sse_json_holds_unicode_line_separators() -> None: + snippet: Final = "Snippet\u2028for\u2029result" + results: Final = [{"title": "Result", "url": "https://example.com/1", "text": snippet}] + body: Final = { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(results, ensure_ascii=False)}]}, + } + + response: Final = _transform_text(f"event: message\ndata: {json.dumps(body, ensure_ascii=False)}\n\n") + + assert _as_tuples(response) == [("Result", "https://example.com/1", snippet, None, None)] diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index a6f7d613f28..635363e77ea 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -587,6 +587,37 @@ class TestChatGPTResponsesMultiAccount: assert headers_default["ChatGPT-Account-Id"] == "acct-default" +def test_chatgpt_non_stream_sse_keeps_an_output_item_whose_text_holds_unicode_line_separators() -> None: + text: Final = "Hello\u2028from\u2029stream" + output_item_done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}, + } + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [], + }, + } + sse_body: Final = "".join( + f"data: {json.dumps(event, ensure_ascii=False)}\n" for event in (output_item_done, completed) + ) + + parsed: Final = ChatGPTResponsesAPIConfig().transform_response_api_response( + model="chatgpt/gpt-6.1-sol", + raw_response=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse_body), + logging_obj=MagicMock(), + ) + + assert parsed.output_text == text + + @pytest.mark.parametrize("call", ["responses", "completion"]) def test_prompt_cache_key_keeps_one_session_header_across_requests( call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py index 3a45c83556c..ccbfc02f0e6 100644 --- a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py +++ b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from typing import Final from unittest.mock import MagicMock @@ -380,6 +380,35 @@ def test_completion_cloudflare_stream(sync_mode: bool, _cloudflare_httpx_transpo assert "language" in content.lower() +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_cloudflare_stream_keeps_unicode_line_separators_inside_content( + sync_mode: bool, _cloudflare_httpx_transport: None +) -> None: + streamed_text: Final = ("alpha\u2028", "beta\u2029", "gamma") + base: Final = {"id": "chatcmpl-cf", "object": "chat.completion.chunk", "created": 1, "model": "@cf/glm"} + body: Final = "".join( + f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': {'content': text}}]}, ensure_ascii=False)}\n\n" + for text in streamed_text + ) + with respx.mock(assert_all_called=True) as api: + api.post(f"{FAKE_API_BASE}/chat/completions").respond( + 200, headers={"content-type": "text/event-stream"}, content=(body + "data: [DONE]\n\n").encode() + ) + request: Final = { + "model": "cloudflare/@cf/glm", + "messages": [{"role": "user", "content": "repeat"}], + "stream": True, + "api_base": FAKE_API_BASE, + "api_key": FAKE_API_KEY, + } + chunks: Final = ( + tuple(completion(**request)) if sync_mode else tuple([c async for c in await acompletion(**request)]) + ) + + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "".join(streamed_text) + + def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: mock = MagicMock(spec=httpx.Response) mock.status_code = 200 @@ -397,13 +426,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch): if sync_mode: - def _iter_lines(): + def _iter_text() -> Iterator[str]: for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" + yield f"data: {chunk}\n\n" + yield "data: [DONE]\n\n" mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() + mock_resp.iter_text.return_value = _iter_text() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} @@ -420,13 +449,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch): mock_post.assert_called_once() else: - async def _aiter_lines(): + async def _aiter_text() -> AsyncIterator[str]: for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" + yield f"data: {chunk}\n\n" + yield "data: [DONE]\n\n" mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() + mock_resp.aiter_text.return_value = _aiter_text() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} diff --git a/tests/unit/llms/codestral/__init__.py b/tests/unit/llms/codestral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/__init__.py b/tests/unit/llms/codestral/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py new file mode 100644 index 00000000000..9ce50d8ceca --- /dev/null +++ b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py @@ -0,0 +1,50 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _fim_stream_body(text: str) -> bytes: + chunks: Final = ( + {"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": None}]}, + {"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}]}, + ) + events: Final = "".join(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" for chunk in chunks) + return f"{events}data: [DONE]\n\n".encode() + + +async def _streamed_texts(sync_stream: bool) -> list[str | None]: + request: Final = { + "model": "text-completion-codestral/codestral-latest", + "prompt": "def f(", + "stream": True, + "api_key": "codestral-key", + "max_retries": 0, + } + if sync_stream: + return [chunk.choices[0].text for chunk in litellm.text_completion(**request)] + return [chunk.choices[0].text async for chunk in await litellm.atext_completion(**request)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta( + sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://codestral.mistral.ai/v1/fim/completions").mock( + return_value=httpx.Response(200, content=_fim_stream_body(_LINE_SEPARATED_TEXT)) + ) + + texts: Final = await _streamed_texts(sync_stream) + + assert texts == [_LINE_SEPARATED_TEXT, None] + assert json.loads(route.calls.last.request.content)["stream"] is True diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py index 83c38b65beb..bb66a70da84 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py @@ -1,11 +1,17 @@ +import json +from typing import Final from unittest.mock import AsyncMock, Mock, patch import aiohttp +import httpx import pytest import litellm +from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ModelResponseStream class TestBaseLLMAIOHTTPHandler: @@ -445,3 +451,34 @@ class TestBaseLLMAIOHTTPHandler: # Should use transport, not connector mock_transport.get_valid_client_session.assert_called_once() assert result is mock_session_from_transport + + +def test_make_sync_call_stream_keeps_unicode_line_separators_inside_a_delta() -> None: + text: Final = "before\u2028between\u2029after" + chunk: Final = { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.5", + "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}], + } + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n".encode() + transport: Final = httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body))) + + completion_stream, _ = BaseLLMAIOHTTPHandler().make_sync_call( + provider_config=AiohttpOpenAIChatConfig(), + api_base="https://example.com/v1/chat/completions", + headers={}, + data={"model": "gpt-5.5", "messages": []}, + model="gpt-5.5", + messages=[], + logging_obj=Mock(), + litellm_params={}, + timeout=5.0, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + contents: Final = [ + chunk.choices[0].delta.content for chunk in completion_stream if isinstance(chunk, ModelResponseStream) + ] + assert contents == [text] diff --git a/tests/unit/llms/custom_httpx/test_sse_lines.py b/tests/unit/llms/custom_httpx/test_sse_lines.py new file mode 100644 index 00000000000..9bb03f6da98 --- /dev/null +++ b/tests/unit/llms/custom_httpx/test_sse_lines.py @@ -0,0 +1,64 @@ +from collections.abc import AsyncIterator, Sequence +from typing import Final + +import pytest + +from litellm.llms.custom_httpx.sse_lines import ( + aiter_sse_lines, + iter_sse_lines, + split_sse_lines, + split_sse_lines_keepends, +) + +CASES: Final = ( + pytest.param(("data: a\u2028b\r\n\r\n",), ("data: a\u2028b", ""), id="u2028-stays-inside-the-line"), + pytest.param( + ("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e\n",), + ("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e",), + id="no-splitlines-only-separator-breaks-a-line", + ), + pytest.param(("a\r\nb\nc\rd\n",), ("a", "b", "c", "d"), id="crlf-lf-and-cr-each-end-a-line"), + pytest.param(("a\r", "\nb\n"), ("a", "b"), id="crlf-split-across-chunks-is-one-break"), + pytest.param(("a\r", "\n\r\n"), ("a", ""), id="crlf-split-then-blank-line"), + pytest.param(("a\r", "", "\nb\n"), ("a", "b"), id="empty-chunk-between-cr-and-lf"), + pytest.param(("a\r", "\r"), ("a", ""), id="cr-then-cr-is-two-breaks"), + pytest.param(("a\r\r\nb\n",), ("a", "", "b"), id="cr-then-crlf-in-one-chunk-is-two-breaks"), + pytest.param(("a\r\nb\r\n\r\n",), ("a", "b", ""), id="crlf-only-chunk"), + pytest.param(("da", "ta: x", "y\n"), ("data: xy",), id="line-spanning-chunks-is-joined"), + pytest.param(("data: x",), ("data: x",), id="unterminated-last-line-is-kept"), + pytest.param(("",), (), id="empty-body"), +) + + +async def _async_chunks(chunks: Sequence[str]) -> AsyncIterator[str]: + for chunk in chunks: + yield chunk + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +def test_iter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None: + assert tuple(iter_sse_lines(chunks)) == tuple(expected) + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +async def test_aiter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None: + assert tuple([line async for line in aiter_sse_lines(_async_chunks(chunks))]) == tuple(expected) + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +def test_split_sse_lines_splits_a_whole_body_only_on_cr_lf_and_crlf( + chunks: Sequence[str], expected: Sequence[str] +) -> None: + assert split_sse_lines("".join(chunks)) == tuple(expected) + + +@pytest.mark.parametrize( + ("text", "expected"), + ( + pytest.param("a\r\nb\rc\nd", ("a\r\n", "b\r", "c\n", "d"), id="each-break-stays-on-its-line"), + pytest.param("data: a\u2028b\u2029c\n\n", ("data: a\u2028b\u2029c\n", "\n"), id="u2028-and-u2029-stay-inside"), + pytest.param("", (), id="empty-body"), + ), +) +def test_split_sse_lines_keepends_splits_only_on_cr_lf_and_crlf(text: str, expected: Sequence[str]) -> None: + assert split_sse_lines_keepends(text) == tuple(expected) diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index b3a224209e9..eb2e1d8e347 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,6 +1,6 @@ import copy import json -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from typing import Final from unittest.mock import MagicMock, patch @@ -2160,13 +2160,13 @@ def mock_chat_response_nonanthropic_prompt_caching() -> Dict[str, Any]: def mock_http_handler_chat_streaming_response() -> MagicMock: mock_stream_chunks = mock_chat_streaming_response_chunks() - def mock_iter_lines(): + def mock_iter_text() -> Iterator[str]: for chunk in mock_stream_chunks: for line in chunk.splitlines(): - yield line + yield f"{line}\n" mock_response = MagicMock() - mock_response.iter_lines.side_effect = mock_iter_lines + mock_response.iter_text.side_effect = mock_iter_text mock_response.status_code = 200 return mock_response @@ -2175,13 +2175,13 @@ def mock_http_handler_chat_streaming_response() -> MagicMock: def mock_http_handler_chat_async_streaming_response() -> MagicMock: mock_stream_chunks = mock_chat_streaming_response_chunks() - async def mock_iter_lines(): + async def mock_aiter_text() -> AsyncIterator[str]: for chunk in mock_stream_chunks: for line in chunk.splitlines(): - yield line + yield f"{line}\n" mock_response = MagicMock() - mock_response.aiter_lines.return_value = mock_iter_lines() + mock_response.aiter_text.return_value = mock_aiter_text() mock_response.status_code = 200 return mock_response diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 04b791cf289..61e5545b8bf 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1329,7 +1329,7 @@ def test_streaming_surfaces_fireworks_response_fields(): raw_response = MagicMock() raw_response.status_code = 200 raw_response.headers = {} - raw_response.iter_lines = lambda: iter(sse_lines) + raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines) client = HTTPHandler() with patch.object(client, "post", return_value=raw_response): @@ -1788,7 +1788,7 @@ def test_streaming_preserves_selected_model_for_private_accounting(): raw_response = MagicMock() raw_response.status_code = 200 raw_response.headers = {} - raw_response.iter_lines = lambda: iter(sse_lines) + raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines) client = HTTPHandler() with patch.object(client, "post", return_value=raw_response): @@ -2275,7 +2275,7 @@ def test_completion_forwards_litellm_user_id_as_user(call_kwargs: dict[str, obje def test_completion_forwards_litellm_user_id_when_streaming() -> None: client: Final = _fireworks_chat_client() - client.post.return_value.iter_lines = lambda: iter(()) + client.post.return_value.iter_text = lambda: iter(()) litellm.completion( model="fireworks_ai/accounts/fireworks/models/kimi-k3", messages=[{"role": "user", "content": "hi"}], diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py index 4119ce99423..53ddb1135b4 100644 --- a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py +++ b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py @@ -4,6 +4,7 @@ Tests for the Google GenAI generateContent guardrail translation handler. import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -190,6 +191,21 @@ async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames(): assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] +@pytest.mark.asyncio +async def test_streaming_raw_sse_text_holding_unicode_line_separators_reaches_the_guardrail_whole() -> None: + guardrail: Final = _mock_guardrail(["clean"]) + text: Final = "harmful\u2028answer\u2029here" + payload: Final = {"candidates": [{"content": {"parts": [{"text": text}]}}]} + chunk: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\r\n\r\n" + + await GoogleGenAIGenerateContentHandler().process_output_streaming_response( + responses_so_far=[chunk], + guardrail_to_apply=guardrail, + ) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [text] + + @pytest.mark.asyncio async def test_streaming_blocking_guardrail_exception_propagates(): handler = GoogleGenAIGenerateContentHandler() diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py new file mode 100644 index 00000000000..06485362162 --- /dev/null +++ b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py @@ -0,0 +1,32 @@ +import json +from typing import Final + +import httpx +import pytest + +from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _messages_event_stream(text: str) -> httpx.Response: + event: Final = ["messages", [[{"type": "AIMessageChunk", "content": text}, {}]]] + return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode())) + + +async def _streamed_contents(iterator: LangGraphSSEStreamIterator, sync_stream: bool) -> list[str | None]: + if sync_stream: + return [chunk.choices[0].delta.content for chunk in iterator] + return [chunk.choices[0].delta.content async for chunk in iterator] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_a_message_chunk(sync_stream: bool) -> None: + iterator: Final = LangGraphSSEStreamIterator( + response=_messages_event_stream(_LINE_SEPARATED_TEXT), model="langgraph/agent" + ) + + contents: Final = await _streamed_contents(iterator, sync_stream) + + assert contents == [_LINE_SEPARATED_TEXT, None] diff --git a/tests/unit/llms/oobabooga/chat/test_oobabooga.py b/tests/unit/llms/oobabooga/chat/test_oobabooga.py index 395a4fb5715..332b6913890 100644 --- a/tests/unit/llms/oobabooga/chat/test_oobabooga.py +++ b/tests/unit/llms/oobabooga/chat/test_oobabooga.py @@ -1,7 +1,12 @@ +import json +from typing import Final from unittest.mock import MagicMock, patch +import httpx +import respx import litellm +from litellm.llms.oobabooga.chat import oobabooga MOCK_COMPLETION_RESPONSE = { "choices": [{"message": {"role": "assistant", "content": "hi there"}}], @@ -50,3 +55,27 @@ def test_url_valued_model_still_targets_that_url(): mock_post.assert_called_once() called_url = mock_post.call_args[0][0] assert called_url == "https://sdk-user.example/v1/chat/completions" + + +def test_stream_keeps_unicode_line_separators_inside_a_data_line(respx_mock: respx.MockRouter) -> None: + text: Final = "before\u2028between\u2029after" + chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}]} + data_line: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}" + respx_mock.post("https://ooba.example/v1/chat/completions").mock( + return_value=httpx.Response(200, content=f"{data_line}\n\ndata: [DONE]\n\n".encode()) + ) + + lines: Final = oobabooga.completion( + model="my-model", + messages=[{"role": "user", "content": "hello"}], + api_base="https://ooba.example", + model_response=litellm.ModelResponse(), + print_verbose=print, + encoding=None, + api_key=None, + logging_obj=MagicMock(), + optional_params={"stream": True}, + litellm_params={}, + ) + + assert [line for line in lines if line] == [data_line, "data: [DONE]"] diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py index 75b60b2dd6c..3df1efe311b 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -4,9 +4,10 @@ from typing import Final from unittest.mock import Mock import httpx +import pytest -from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.openai_like.chat.handler import make_sync_call +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai_like.chat.handler import make_call, make_sync_call _FAKE_STREAMED_REPLY: Final = { "id": "chatcmpl-fake-stream", @@ -73,3 +74,35 @@ def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail(): assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"] assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64 assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920 + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _streamed_reply(request: httpx.Request) -> httpx.Response: + chunk: Final = {"id": "chatcmpl-separators", "choices": [{"index": 0, "delta": {"content": _LINE_SEPARATED_TEXT}}]} + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n" + return httpx.Response(200, stream=httpx.ByteStream(body.encode())) + + +async def _streamed_chunks(sync_stream: bool) -> list[dict[str, object]]: + call_args: Final = { + "api_base": "https://openai-like.example/v1/chat/completions", + "headers": {}, + "data": json.dumps({"model": "openai/gpt-oss-120b", "messages": [], "stream": True}), + "model": "openai/gpt-oss-120b", + "messages": [], + "logging_obj": Mock(), + } + transport: Final = httpx.MockTransport(_streamed_reply) + if sync_stream: + return list(make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args)) + return [chunk async for chunk in await make_call(client=AsyncHTTPHandler(transport=transport), **call_args)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_the_delta(sync_stream: bool) -> None: + chunks: Final = await _streamed_chunks(sync_stream) + + assert "".join(str(chunk["text"]) for chunk in chunks) == _LINE_SEPARATED_TEXT diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py index f9d12b7eeea..36c319d9524 100644 --- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -902,20 +902,14 @@ class TestSnowflakeChatCompletion: @pytest.mark.parametrize("sync_mode", [True, False]) def test_chat_completion_snowflake_stream(self, sync_mode): - raw_chunks = _streaming_chunks() + sse_body = "".join(f"data: {chunk}\n\n" for chunk in _streaming_chunks()) + "data: [DONE]\n\n" + mock_resp = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=httpx.ByteStream(sse_body.encode()), + ) if sync_mode: - - def _iter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="snowflake/mistral-7b", @@ -930,16 +924,6 @@ class TestSnowflakeChatCompletion: mock_post.assert_called_once() else: - async def _aiter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - async def _run(): with patch.object( AsyncHTTPHandler, @@ -964,8 +948,8 @@ class TestSnowflakeChatCompletion: chunks_received = asyncio.run(_run()) - assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + assert content == "The sky is blue" async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index a02840ed4ef..7afbd05b221 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -1,3 +1,4 @@ +import json from typing import Final from unittest.mock import AsyncMock, Mock @@ -639,3 +640,49 @@ def test_predibase_completion_passes_existing_config_to_async_completion(monkeyp assert result == "async-result" assert captured["async_kwargs"]["predibase_config"] is captured["config_instance"] + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _generate_stream_body(text: str) -> bytes: + events: Final = ( + {"token": {"text": text}}, + {"token": {"text": ""}, "generated_text": text, "details": {"finish_reason": "eos_token"}}, + ) + return "".join(f"data:{json.dumps(event, ensure_ascii=False)}\n\n" for event in events).encode() + + +async def _streamed_contents(sync_stream: bool) -> list[str | None]: + request: Final = { + "model": "predibase/llama-3-8b-instruct", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "predibase_tenant_id": "tenant-1", + "api_key": "pb-key", + "max_retries": 0, + } + if sync_stream: + return [chunk.choices[0].delta.content for chunk in litellm.completion(**request)] + return [chunk.choices[0].delta.content async for chunk in await litellm.acompletion(**request)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_predibase_stream_keeps_unicode_line_separators_inside_a_token( + sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "custom_prompt_dict", + {"llama-3-8b-instruct": {"roles": {}, "initial_prompt_value": "", "final_prompt_value": ""}}, + ) + respx_mock.post( + "https://serving.app.predibase.com/tenant-1/deployments/v2/llms/llama-3-8b-instruct/generate_stream" + ).mock(return_value=httpx.Response(200, content=_generate_stream_body(_LINE_SEPARATED_TEXT))) + + contents: Final = await _streamed_contents(sync_stream) + + assert contents == [_LINE_SEPARATED_TEXT, None] diff --git a/tests/unit/llms/triton/test_triton.py b/tests/unit/llms/triton/test_triton.py index 7c0de603a3e..238c30816b2 100644 --- a/tests/unit/llms/triton/test_triton.py +++ b/tests/unit/llms/triton/test_triton.py @@ -1,5 +1,7 @@ import json import traceback +from collections.abc import Iterator +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -210,8 +212,8 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock): mock_response = MagicMock() if stream: - def mock_iter_lines(): - mock_output = "".join( + def mock_iter_text() -> Iterator[str]: + mock_output: Final = "".join( [ 'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"' + t @@ -219,10 +221,9 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock): for t in ["I", " am", " an", " AI", " assistant"] ] ) - for out in mock_output.split("\n"): - yield out + yield mock_output - mock_response.iter_lines = mock_iter_lines + mock_response.iter_text = mock_iter_text else: def return_val(): diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py index 06aaca44e8a..20333fd0465 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -14,7 +14,7 @@ import pytest from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, get_async_httpx_client from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) @@ -170,3 +170,48 @@ async def test_async_stream_wrapper_without_a_client_posts_through_the_cached_ve assert type(wrapper) is CustomStreamWrapper assert [str(request.url) for request in sent_requests] == [api_base] assert json.loads(sent_requests[0].content) == {"class_method": "stream_query"} + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _separated_text_reply(request: httpx.Request) -> httpx.Response: + chunk: Final = {"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}], "role": "model"}, "finish_reason": "STOP"} + return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode())) + + +async def _streamed_text(sync_stream: bool) -> str: + messages: Final = [{"role": "user", "content": "hi"}] + call_args: Final = { + "model": "agent_engine/123", + "custom_llm_provider": "vertex_ai", + "logging_obj": Logging( + model="agent_engine/123", + messages=messages, + stream=True, + call_type="completion" if sync_stream else "acompletion", + start_time=datetime(2026, 1, 1), + litellm_call_id="call-separators", + function_id="fn-separators", + ), + "api_base": "https://us-central1-aiplatform.googleapis.com/v1/reasoningEngines/123:streamQuery", + "headers": {}, + "data": {"class_method": "stream_query"}, + "messages": messages, + "litellm_params": {}, + } + transport: Final = httpx.MockTransport(_separated_text_reply) + if sync_stream: + sync_wrapper: Final = VertexAgentEngineConfig().get_sync_custom_stream_wrapper( + client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args + ) + return "".join(chunk.choices[0].delta.content or "" for chunk in sync_wrapper) + async_wrapper: Final = await VertexAgentEngineConfig().get_async_custom_stream_wrapper( + client=AsyncHTTPHandler(transport=transport), **call_args + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in async_wrapper]) + + +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_the_text(sync_stream: bool) -> None: + assert await _streamed_text(sync_stream) == _LINE_SEPARATED_TEXT diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 3e1bca18f08..67c78b62151 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3,8 +3,9 @@ import json import logging import re import threading -from collections.abc import Awaitable, Callable, Iterator, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from copy import deepcopy +from itertools import pairwise from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch @@ -34,7 +35,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ) from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage +from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger @@ -10125,7 +10126,7 @@ def test_vertex_ai_streaming_response_id(): mock_response = MagicMock() mock_response.status_code = 200 - mock_response.iter_lines = MagicMock(return_value=stream_response()) + mock_response.iter_text = MagicMock(return_value=stream_response()) return mock_response logging_obj = MagicMock() @@ -10690,6 +10691,63 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n assert sent["includeThoughts"] is expected +class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._chunks: Final = chunks + + def __iter__(self) -> Iterator[bytes]: + yield from self._chunks + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + +def _split_at(body: bytes, cuts: tuple[int, ...]) -> tuple[bytes, ...]: + return tuple(body[start:end] for start, end in pairwise((0, *sorted(cuts), len(body)))) + + +async def _gemini_stream(is_async: bool, chunks: tuple[bytes, ...]) -> list[ModelResponseStream]: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_BodyInChunks(chunks), headers={"content-type": "text/event-stream"}) + + kwargs: Final = { + "model": "gemini/gemini-3.8-flash", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "stream_options": {"include_usage": True}, + "api_key": "fake-gemini-key", + } + if is_async: + stream: Final = await litellm.acompletion( + **kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(handler)) + ) + return [chunk async for chunk in stream] + return list( + litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))) + ) + + +@pytest.mark.parametrize("is_async", [True, False]) +async def test_gemini_stream_keeps_text_and_usage_when_a_part_contains_a_unicode_line_separator(is_async: bool) -> None: + events: Final = ( + {"candidates": [{"content": {"role": "model", "parts": [{"text": "line one\u2028line two"}]}}]}, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": " and three"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18}, + }, + ) + body: Final = "".join(f"data: {json.dumps(event, ensure_ascii=False)}\r\n\r\n" for event in events).encode() + cut_inside_separator_and_crlf: Final = (body.index(b"\xe2\x80\xa8") + 1, body.index(b"\r") + 1) + + received: Final = await _gemini_stream(is_async, _split_at(body, cut_inside_separator_and_crlf)) + + usage: Final = next(chunk.usage for chunk in reversed(received) if getattr(chunk, "usage", None)) + text: Final = "".join(c.choices[0].delta.content or "" for c in received if c.choices) + assert text == "line one\u2028line two and three" + assert (usage.prompt_tokens, usage.completion_tokens) == (10, 8) + + @pytest.mark.parametrize("chunk_indices", [[0, 1], [2, 1, 0]]) def test_grounding_support_preserves_every_web_source(chunk_indices: list[int]) -> None: metadata: Final = { diff --git a/tests/unit/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py index 34502d5cebe..0ed469acfa7 100644 --- a/tests/unit/llms/vertex_ai/test_http_status_201.py +++ b/tests/unit/llms/vertex_ai/test_http_status_201.py @@ -39,8 +39,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): # Create a mock response with status code 201 mock_response = MagicMock() mock_response.status_code = 201 - mock_response.aiter_lines = MagicMock() - mock_response.aiter_lines.return_value = ["test response"] + mock_response.aiter_text = MagicMock() + mock_response.aiter_text.return_value = ["test response"] # Setup mock client mock_client = MagicMock() @@ -76,8 +76,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): # Create a mock response with status code 201 mock_response = MagicMock() mock_response.status_code = 201 - mock_response.iter_lines = MagicMock() - mock_response.iter_lines.return_value = ["test response"] + mock_response.iter_text = MagicMock() + mock_response.iter_text.return_value = ["test response"] mock_post.return_value = mock_response # Call the make_sync_call function diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index a7d7ac5627c..7abe6ad7799 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -4,10 +4,11 @@ Mock tests for A2A endpoints. Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request. """ +import asyncio import json import socket import sys -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import AbstractContextManager, ExitStack from dataclasses import dataclass from typing import Final @@ -24,6 +25,21 @@ from litellm.types.agents import AgentCaller AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] +class _SSELineStream(httpx.AsyncByteStream): + def __init__(self, lines: tuple[str, ...], first_line_delay: float) -> None: + self._lines = lines + self._first_line_delay = first_line_delay + + async def __aiter__(self) -> AsyncIterator[bytes]: + await asyncio.sleep(self._first_line_delay) + for line in self._lines: + yield f"{line}\n".encode() + + +def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx.Response: + return httpx.Response(200, stream=_SSELineStream(lines, first_line_delay)) + + @dataclass(frozen=True, slots=True) class CapturedAgentCall: request_id: object @@ -941,13 +957,7 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str mock_http_response.json.return_value = upstream_response mock_http_response.is_success = True - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}') mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -982,8 +992,8 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str user_api_key_dict=user_api_key_dict, ) if method == "tasks/resubscribe": - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "task-1" in streamed if method == "tasks/resubscribe": forwarded_body = mock_async_client.build_request.call_args.kwargs["json"] @@ -1007,14 +1017,7 @@ async def test_subscribe_to_task_returns_sse_stream(): 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}', ] - async def fake_aiter_lines(): - for line in sse_lines: - yield line - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response(*sse_lines) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1064,13 +1067,9 @@ async def test_subscribe_to_task_calls_pre_call_hook(): mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response( + 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}' + ) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1117,9 +1116,9 @@ async def test_subscribe_to_task_calls_pre_call_hook(): ) assert response.media_type == "text/event-stream" - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "completed" in streamed mock_proxy_logging.pre_call_hook.assert_awaited_once() call_kwargs = mock_proxy_logging.pre_call_hook.await_args.kwargs assert call_kwargs.get("call_type") == "asend_message" @@ -1149,16 +1148,10 @@ async def test_subscribe_to_task_runs_post_call_streaming_guardrail(): mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - async def fake_aiter_lines(): - yield ( - 'data: {"jsonrpc":"2.0","id":"req-1","result":' - '{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}' - ) - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response( + 'data: {"jsonrpc":"2.0","id":"req-1","result":' + '{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}' + ) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1600,14 +1593,7 @@ async def test_pascal_method_names_normalize_to_wire_format(pascal_method: str, mock_http_response.is_success = True mock_http_response.raise_for_status = MagicMock() - async def _empty_aiter_lines(): - return - yield # make it an async generator - - mock_sse_resp = AsyncMock() - mock_sse_resp.is_success = True - mock_sse_resp.aiter_lines = _empty_aiter_lines - mock_sse_resp.aclose = AsyncMock() + mock_sse_resp = _sse_upstream_response() mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -2349,13 +2335,7 @@ async def test_task_methods_forward_trace_header(method: str): mock_http_response.json.return_value = upstream_response mock_http_response.is_success = True - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}') mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -2390,8 +2370,8 @@ async def test_task_methods_forward_trace_header(method: str): user_api_key_dict=user_api_key_dict, ) if method == "tasks/resubscribe": - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "task-1" in streamed if method == "tasks/resubscribe": forwarded_headers = mock_async_client.build_request.call_args.kwargs["headers"] @@ -2646,17 +2626,12 @@ def test_served_version_falls_back_to_header_when_unconfigured(): assert _served_version(_agent(None), _request_with_a2a_header(None)) == "0.3" -def _sse_agent_handler(lines): - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = lines - mock_resp.aclose = AsyncMock() - - mock_async_client = MagicMock() +def _sse_agent_handler(upstream_response: httpx.Response) -> MagicMock: + mock_async_client: Final = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) - mock_async_client.send = AsyncMock(return_value=mock_resp) + mock_async_client.send = AsyncMock(return_value=upstream_response) - mock_handler = MagicMock() + mock_handler: Final = MagicMock() mock_handler.client = mock_async_client return mock_handler @@ -2685,13 +2660,13 @@ async def test_forward_jsonrpc_sse_pings_while_the_upstream_agent_is_still_silen monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) - async def _slow_lines(): - await asyncio.sleep(0.3) - yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}' + slow_upstream = _sse_upstream_response( + 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}', first_line_delay=0.3 + ) with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", - return_value=_sse_agent_handler(_slow_lines), + return_value=_sse_agent_handler(slow_upstream), ): response = await _resubscribe_response() assert response.headers["x-accel-buffering"] == "no" @@ -2713,12 +2688,11 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) - async def _lines(): - yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}' - with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", - return_value=_sse_agent_handler(_lines), + return_value=_sse_agent_handler( + _sse_upstream_response('data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}') + ), ): response = await _resubscribe_response() assert "x-accel-buffering" not in response.headers @@ -2728,6 +2702,28 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure assert json.loads(chunks[-1].removeprefix("data: "))["result"]["kind"] == "task" +@pytest.mark.asyncio +async def test_forward_jsonrpc_sse_keeps_unicode_line_separators_inside_an_event(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) + separated_text: Final = "before\u2028between\u2029after" + event: Final = { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "parts": [{"kind": "text", "text": separated_text}]}, + } + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=_sse_agent_handler(_sse_upstream_response(f"data: {json.dumps(event, ensure_ascii=False)}")), + ): + response: Final = await _resubscribe_response() + chunks: Final = [chunk async for chunk in response.body_iterator] + + assert [json.loads(chunk.removeprefix("data: "))["result"]["parts"][0]["text"] for chunk in chunks] == [ + separated_text + ] + + async def _stream_message_response(): from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message diff --git a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py index 85bb8304b95..5911fcdc073 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -4,6 +4,7 @@ Tests for restamping the public model on Anthropic Messages streaming chunks. import json from collections.abc import AsyncIterator, Sequence +from typing import Final import pytest @@ -46,6 +47,28 @@ def test_restamps_sse_bytes_frame(): assert b"event: message_start" in restamped +def test_restamps_sse_bytes_frame_holding_a_line_separator() -> None: + payload: Final = { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5-20251001", + "content": [{"type": "text", "text": "alpha\u2028beta"}], + }, + } + frame: Final = f"event: message_start\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n".encode() + + restamped: Final = restamp_anthropic_stream_chunk_model(frame, "claude-auto-1") + + assert isinstance(restamped, bytes) + data_line: Final = next(line for line in restamped.decode().split("\n") if line.startswith("data:")) + message: Final = json.loads(data_line[len("data:") :])["message"] + assert message["model"] == "claude-auto-1" + assert message["content"] == [{"type": "text", "text": "alpha\u2028beta"}] + + def test_restamps_event_dict(): chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index e09a179eadf..47b6af9aeb3 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2503,6 +2503,33 @@ def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects( assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is interrupted + +def test_stream_was_interrupted_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"}, + "usage": {"output_tokens": 3}, + } + all_chunks: Final = [ + 'data: {"type": "content_block_stop", "index": 0}', + f"data: {json.dumps(message_delta, ensure_ascii=False)}", + ] + + assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is False + + +def test_extract_model_reads_a_message_start_whose_json_holds_unicode_line_separators() -> None: + message_start: Final = { + "type": "message_start", + "message": { + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Prefill\u2028line\u2029two"}], + }, + } + chunk: Final = f"event: message_start\ndata: {json.dumps(message_start, ensure_ascii=False)}\n\n" + + assert AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks([chunk]) == "claude-sonnet-5" + @pytest.mark.asyncio async def test_anthropic_passthrough_handler(mock_httpx_response, mock_response, mock_logging_obj): """ @@ -3004,3 +3031,24 @@ def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_m (entry,) = _recorded_guardrail_entries(mock_logging_obj) assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015) + + +def test_anthropic_passthrough_stream_bills_a_bedrock_guardrail_trace_whose_json_holds_unicode_line_separators( + monkeypatch: pytest.MonkeyPatch, mock_logging_obj: LiteLLMLoggingObj, all_chunks: list[str] +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + trace: Final = {"guardrail": {**_BEDROCK_INVOKE_TRACE["guardrail"], "actionReason": "Blocked\u2028by\u2029policy"}} + stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": trace}, ensure_ascii=False) + AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks( + litellm_logging_obj=mock_logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude", "messages": []}, + endpoint_type=MagicMock(), + start_time=datetime(2025, 1, 1), + all_chunks=[*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"], + end_time=datetime(2025, 1, 1), + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py index 1b83c5140ca..b0f84f2ecbe 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py @@ -1,6 +1,7 @@ import asyncio import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -267,6 +268,12 @@ class TestRunIdFromSseFrames: frames = b'data: {"run_id": "run-7", "event": "INITIALIZED"}\n\ndata: {"run_id": "run-7", "event": "ACTION"}\n\n' assert run_id_from_sse_frames(frames) == "run-7" + def test_finds_run_id_in_a_frame_whose_json_holds_unicode_line_separators(self) -> None: + event: Final = {"run_id": "run-7", "event": "ACTION", "message": "Clicked\u2028Submit\u2029button"} + frames: Final = f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + assert run_id_from_sse_frames(frames) == "run-7" + def test_skips_frames_without_run_id(self): frames = b': keepalive\n\ndata: not-json\n\ndata: {"event": "HEARTBEAT"}\n\n' assert run_id_from_sse_frames(frames) is None diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index 49bc2e9c29e..2496aa47d22 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -30,7 +30,7 @@ from litellm.llms.base_llm.managed_resources.utils import( resolve_passthrough_managed_id_provider, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from typing import Any +from typing import Any, Final def _user() -> UserAPIKeyAuth: @@ -254,6 +254,37 @@ async def test_streamed_response_with_cr_only_frame_delimiters_is_still_owned_an assert output == _response_stream_bytes(managed_id).replace(b"\n", b"\r") +def _instructed_response_stream_bytes(raw_id: str, instructions: str) -> bytes: + events: Final = ( + {"type": "response.created", "response": {"id": raw_id, "status": "in_progress", "instructions": instructions}}, + {"type": "response.completed", "response": {"id": raw_id, "status": "completed", "instructions": instructions}}, + ) + return b"".join( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() for event in events + ) + + +@pytest.mark.asyncio +async def test_streamed_response_whose_json_holds_unicode_line_separators_is_still_owned_and_rewritten() -> None: + instructions: Final = "Answer\u2028in\u2029verse" + pc: Final = _prisma_client() + + output: Final = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(_instructed_response_stream_bytes(RAW_RESPONSE_ID, instructions), 4096), + provider="openai", + method="POST", + route="/openai_passthrough/v1/responses", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + pc.db.litellm_managedobjecttable.upsert.assert_awaited_once() + managed_id: Final = pc.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]["unified_object_id"] + assert output == _instructed_response_stream_bytes(managed_id, instructions) + + @pytest.mark.asyncio async def test_streamed_bytes_untouched_on_routes_without_a_response_id(): pc = _prisma_client() diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 63e38ce533a..42acad28488 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -17,7 +17,7 @@ following the OpenAI Response API format. import asyncio import json from datetime import datetime, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, Mock, patch import pytest @@ -1400,7 +1400,7 @@ def _make_sse_stream(events: list) -> Mock: async def _body_iterator(): for event in events: - yield f"data: {json.dumps(event)}" + yield f"data: {json.dumps(event, ensure_ascii=False)}" yield "data: [DONE]" mock_response = Mock() @@ -1639,6 +1639,33 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "completed" assert final_call.kwargs["usage"] == {"input_tokens": 10, "output_tokens": 50} + @pytest.mark.asyncio + async def test_response_completed_whose_json_holds_unicode_line_separators_keeps_its_output(self) -> None: + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + item: Final = { + "id": "item_1", + "type": "message", + "content": [{"type": "output_text", "text": "Line\u2028one\u2029two"}], + } + events: Final = [ + {"type": "response.completed", "response": {"id": "resp_123", "status": "completed", "output": [item]}} + ] + handler: Final = AsyncMock(spec=ResponsePollingHandler) + + with patch( + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=_make_sse_stream(events)) + await background_streaming_task(**_make_background_streaming_kwargs("poll_u2028", handler)) + + calls: Final = handler.update_state.call_args_list + outputs: Final = [call.kwargs["output"] for call in calls if "output" in call.kwargs] + assert outputs[-1] == [item] + assert calls[-1].kwargs["status"] == "completed" + @pytest.mark.asyncio async def test_fallback_status_derived_from_event_type_when_status_field_missing( self, diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index ed90acbd2fa..0b971bc1696 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1288,6 +1288,106 @@ async def test_reasoning_item_closes_before_a_tool_call_from_the_same_chunk_open ] +def _text_opening_a_tool_call_chunk(content: str, arguments: str = "") -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=content, + tool_calls=[ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": arguments}, + "index": 0, + } + ], + ), + finish_reason=None, + ) + ], + ) + + +def _tool_arguments_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[{"type": "function", "function": {"arguments": arguments}, "index": 0}], + ), + finish_reason=finish_reason, + ) + ], + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_call_opened_in_a_text_chunk_streams_each_argument_delta(sync_mode: bool) -> None: + argument_pieces: Final = ('{"', "city", '"', ': "', "Paris", '"', "}") + iterator: Final = _build_iterator( + [ + _text_opening_a_tool_call_chunk(""), + *(_tool_arguments_chunk(piece) for piece in argument_pieces), + _chunk("", finish_reason="tool_calls"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "" + assert [ + event.delta + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] == list(argument_pieces) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_call_sharing_a_chunk_with_text_opens_after_the_whole_message(sync_mode: bool) -> None: + iterator: Final = _build_iterator( + [ + _text_opening_a_tool_call_chunk("Checking ", arguments='{"city": "Paris"}'), + _chunk("the "), + _chunk("weather now."), + _chunk("", finish_reason="tool_calls"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + event_types: Final = [getattr(event, "type", None) for event in events] + function_call_added_at: Final = next( + position + for position, event in enumerate(events) + if event_types[position] == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and event.item.type == "function_call" + ) + streamed: Final = [ + event + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) + ] + assert ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA not in event_types[function_call_added_at:] + assert [event.sequence_number for event in streamed] == sorted(event.sequence_number for event in streamed) + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Checking the weather now." + assert _event_text(events, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) == '{"city": "Paris"}' + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_tool_only_stream_puts_the_function_call_at_output_index_zero(sync_mode: bool): diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py index e01b9120416..e58392ca266 100644 --- a/tests/unit/sandbox/test_e2b_sandbox.py +++ b/tests/unit/sandbox/test_e2b_sandbox.py @@ -7,6 +7,7 @@ integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. """ import json +from typing import Final import httpx import pytest @@ -19,19 +20,18 @@ from litellm.llms.e2b.sandbox.transformation import ( ) +def _line_stream(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) + + class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): + def __init__(self, *, json_data: object = None, status_code: int = 200) -> None: self._json = json_data - self._lines = lines or [] self.status_code = status_code def json(self): return self._json - async def aiter_lines(self): - for line in self._lines: - yield line - class FakeHTTPClient: """Records outbound requests and returns canned responses keyed by URL.""" @@ -61,7 +61,7 @@ class FakeHTTPClient: if url.endswith("/execute"): if self.execute_raises is not None: raise self.execute_raises - return FakeResponse(lines=self.execute_lines) + return _line_stream(self.execute_lines) raise AssertionError(f"unexpected POST {url}") async def delete(self, url, headers=None, **kwargs): @@ -249,6 +249,20 @@ def test_parse_lines_skips_non_json_lines(): assert result.error is None +@pytest.mark.asyncio +async def test_run_code_keeps_unicode_line_separators_inside_stdout() -> None: + separated_text: Final = "before\u2028between\u2029after\n" + client: Final = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": separated_text, "timestamp": 1}, ensure_ascii=False)] + ) + handle: Final = ContainerHandle(id="sbx_sep", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok"} + + result: Final = await E2BSandboxConfig().arun_code(container=handle, code="print(text)", client=client) + + assert result.stdout == separated_text + + @pytest.mark.asyncio async def test_run_code_aborts_on_output_over_cap(): big_line = "x" * (MAX_OUTPUT_BYTES + 1) diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py index 2928dea100e..36ae848786c 100644 --- a/tests/unit/sandbox/test_opensandbox_sandbox.py +++ b/tests/unit/sandbox/test_opensandbox_sandbox.py @@ -27,10 +27,13 @@ def sse(data): return f"data: {json.dumps(data)}" +def _line_stream(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) + + class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): + def __init__(self, *, json_data: object = None, status_code: int = 200) -> None: self._json = json_data - self._lines = lines or [] self.status_code = status_code def json(self): @@ -40,10 +43,6 @@ class FakeResponse: if self.status_code >= 400: raise http_status_error(self.status_code) - async def aiter_lines(self): - for line in self._lines: - yield line - class FakeHTTPClient: def __init__( @@ -93,7 +92,7 @@ class FakeHTTPClient: if url.endswith("/code"): if self.execute_raises is not None: raise self.execute_raises - return FakeResponse(lines=self.execute_lines) + return _line_stream(self.execute_lines) raise AssertionError(f"unexpected POST {url}") async def get(self, url, headers=None, params=None, **kwargs): diff --git a/tests/unit/test_groq_streaming_encoding.py b/tests/unit/test_groq_streaming_encoding.py index ffe87fb9890..37f05272f67 100644 --- a/tests/unit/test_groq_streaming_encoding.py +++ b/tests/unit/test_groq_streaming_encoding.py @@ -6,148 +6,58 @@ UTF-8 encoded content in streaming responses, specifically fixing the ASCII encoding error described in issue #12660. """ -import asyncio -from unittest.mock import AsyncMock, Mock +import json +from typing import Final +from unittest.mock import Mock +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai_like.chat.handler import make_call, make_sync_call - -class MockResponse: - """Mock httpx response for testing UTF-8 handling.""" - - def __init__(self, test_content: str): - self.test_content = test_content - self.status_code = 200 - - def iter_text(self, encoding="utf-8"): - """Mock iter_text that yields content with the specified encoding.""" - yield self.test_content - - async def aiter_text(self, encoding="utf-8"): - """Mock aiter_text that yields content with the specified encoding.""" - yield self.test_content - - def iter_lines(self): - """Mock iter_lines method for synchronous streaming.""" - yield self.test_content - - async def aiter_lines(self): - """Mock aiter_lines method for asynchronous streaming.""" - yield self.test_content - - def json(self): - return {"choices": [{"delta": {"content": "test"}}]} +_UNICODE_TEXTS: Final = ( + "The symbol µ represents micro", + "Testing © character", + "Testing ™ character", + "Testing € character", + "Testing 北京 character", + "Testing 🚀 character", + "Testing Ñoño character", +) -class MockSyncClient: - """Mock synchronous HTTP client for testing.""" - - def __init__(self, response_content: str): - self.response_content = response_content - - def post(self, *args, **kwargs): - return MockResponse(self.response_content) +def _utf8_stream_transport(text: str) -> httpx.MockTransport: + chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}}]} + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode() + return httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body))) -class MockAsyncClient: - """Mock asynchronous HTTP client for testing.""" - - def __init__(self, response_content: str): - self.response_content = response_content - - async def post(self, *args, **kwargs): - return MockResponse(self.response_content) +def _call_args() -> dict[str, object]: + return { + "api_base": "https://test.com/v1/chat/completions", + "headers": {"Authorization": "Bearer test"}, + "data": '{"model": "test", "messages": []}', + "model": "test-model", + "messages": [], + "logging_obj": Mock(), + } -def test_utf8_streaming_sync(): - """Test that synchronous streaming handles UTF-8 characters correctly.""" - # Content with the µ character that was causing issues - test_content = ( - 'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n' +@pytest.mark.parametrize("text", _UNICODE_TEXTS) +def test_utf8_streaming_sync(text: str) -> None: + completion_stream: Final = make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=_utf8_stream_transport(text))), **_call_args() ) - mock_client = MockSyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error - completion_stream = make_sync_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - # Verify we can iterate through the stream without encoding errors - assert completion_stream is not None + assert "".join(chunk["text"] for chunk in completion_stream) == text @pytest.mark.asyncio -async def test_utf8_streaming_async(): - """Test that asynchronous streaming handles UTF-8 characters correctly.""" - # Content with the µ character that was causing issues - test_content = ( - 'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n' +@pytest.mark.parametrize("text", _UNICODE_TEXTS) +async def test_utf8_streaming_async(text: str) -> None: + completion_stream: Final = await make_call( + client=AsyncHTTPHandler(transport=_utf8_stream_transport(text)), **_call_args() ) - mock_client = MockAsyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error - completion_stream = await make_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - # Verify we can iterate through the stream without encoding errors - assert completion_stream is not None - - -def test_various_unicode_characters(): - """Test streaming with various Unicode characters that could cause issues.""" - unicode_test_cases = [ - "µ", # Micro symbol (the original issue) - "©", # Copyright symbol - "™", # Trademark symbol - "€", # Euro symbol - "北京", # Chinese characters - "🚀", # Emoji - "Ñoño", # Spanish characters with tildes - ] - - for unicode_char in unicode_test_cases: - test_content = f'data: {{"choices":[{{"delta":{{"content":"Testing {unicode_char} character"}}}}]}}\n\n' - - mock_client = MockSyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error for any Unicode character - completion_stream = make_sync_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - assert ( - completion_stream is not None - ), f"Failed to handle Unicode character: {unicode_char}" - - -if __name__ == "__main__": - test_utf8_streaming_sync() - asyncio.run(test_utf8_streaming_async()) - test_various_unicode_characters() - print("All UTF-8 streaming tests passed!") + assert "".join([chunk["text"] async for chunk in completion_stream]) == text From 7650f766b006077442abe758ee4b22b5c436ad18 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Wed, 28 Oct 2026 17:01:52 -0700 Subject: [PATCH 3/9] refactor(router): fence the router surface and document the router-core boundary (#48839) Add rustdoc overviews for a new litellm-router-core crate and for litellm-model-catalog, a module docstring on litellm/router.py, and a shrink-only allowlist of the Router members and router_utils modules used outside the router, checked in CI by router_surface_check.py. Move the router-free helpers out of litellm/router_utils: batch_utils to litellm_core_utils, and the auto-router naming contract, tuning baseline and baseline request next to the auto-router plugin. No behaviour change. --- .github/workflows/test-linting.yml | 3 + litellm-rust/Cargo.lock | 4 + litellm-rust/crates/model-catalog/src/lib.rs | 33 +++ litellm-rust/crates/router-core/Cargo.toml | 9 + litellm-rust/crates/router-core/src/lib.rs | 63 +++++ .../batch_utils.py | 0 .../prompt_templates/common_utils.py | 2 +- .../llms/anthropic/prompt_cache_prediction.py | 2 +- litellm/proxy/common_request_processing.py | 2 +- .../guardrails/auto_router_compression.py | 2 +- litellm/proxy/health_check.py | 2 +- .../proxy/hooks/autorouter_baseline_cache.py | 2 +- .../proxy/hooks/model_max_budget_limiter.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 2 +- .../auto_router_endpoints.py | 12 +- .../model_management_endpoints.py | 31 +-- .../auto_router_availability.py | 4 +- .../auto_router_permissions.py | 5 +- litellm/proxy/proxy_server.py | 14 +- litellm/router.py | 61 +++-- .../auto_router}/auto_router_model_naming.py | 0 .../auto_router_tuning_baseline.py | 0 .../auto_router}/baseline_request.py | 0 litellm/router_strategy/budget_limiter.py | 2 +- litellm/router_strategy/least_busy.py | 2 +- litellm/router_strategy/lowest_cost.py | 2 +- litellm/router_strategy/lowest_latency.py | 2 +- litellm/router_strategy/lowest_tpm_rpm.py | 2 +- litellm/router_strategy/lowest_tpm_rpm_v2.py | 2 +- .../router_utils/fallback_event_handlers.py | 2 +- .../router_code_coverage.py | 4 +- .../router_surface_allowlist.json | 211 ++++++++++++++++ .../router_surface_check.py | 235 ++++++++++++++++++ .../unbounded_in_baseline.txt | 2 +- .../test_batch_utils.py | 13 +- .../test_auto_router_endpoints.py | 4 +- .../test_model_management_endpoints.py | 49 ++-- .../test_auto_router_availability.py | 7 +- .../unit/proxy/proxy_server/test_lifecycle.py | 12 +- .../test_auto_router_model_naming.py | 8 +- .../test_auto_router_tuning_baseline.py | 4 +- .../test_baseline_request.py | 2 +- .../router_strategy/test_complexity_router.py | 36 ++- tests/unit/router_strategy/test_llm_v2.py | 4 +- tests/unit/test_private_usage_aliases.py | 2 +- .../test_router/test_router_helper_utils.py | 15 +- .../test_router/test_router_surface_check.py | 58 +++++ 47 files changed, 797 insertions(+), 138 deletions(-) create mode 100644 litellm-rust/crates/router-core/Cargo.toml create mode 100644 litellm-rust/crates/router-core/src/lib.rs rename litellm/{router_utils => litellm_core_utils}/batch_utils.py (100%) rename litellm/{router_utils => router_strategy/auto_router}/auto_router_model_naming.py (100%) rename litellm/{router_utils => router_strategy/auto_router}/auto_router_tuning_baseline.py (100%) rename litellm/{router_utils => router_strategy/auto_router}/baseline_request.py (100%) create mode 100644 tests/code_coverage_tests/router_surface_allowlist.json create mode 100644 tests/code_coverage_tests/router_surface_check.py rename tests/unit/{router_utils => litellm_core_utils}/test_batch_utils.py (97%) rename tests/unit/{router_utils => router_strategy}/test_auto_router_model_naming.py (99%) rename tests/unit/{router_utils => router_strategy}/test_auto_router_tuning_baseline.py (99%) rename tests/unit/{router_utils => router_strategy}/test_baseline_request.py (95%) create mode 100644 tests/unit/test_router/test_router_surface_check.py diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index b9f8d7e7dfb..89ca60ad9d1 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -561,6 +561,9 @@ jobs: - name: router_code_coverage run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py + - name: router_surface_check + run: uv run --no-sync python ./tests/code_coverage_tests/router_surface_check.py + - name: test_chat_completion_imports run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9c45c5c4b1c..10e590b5006 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4393,6 +4393,10 @@ dependencies = [ "rstest", ] +[[package]] +name = "litellm-router-core" +version = "0.1.0" + [[package]] name = "litellm-router-types" version = "0.1.0" diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index 066c2c83c6b..d42c05e6fe7 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -1,3 +1,36 @@ +//! The model catalog: per-model pricing, context limits and capabilities. +//! +//! This is the Rust form of `model_prices_and_context_window.json`, the map Python reads as +//! `litellm.model_cost` and through `litellm.get_model_info`. It answers "what does this model +//! cost and what can it do"; it knows nothing about configured deployments. The router's view of +//! deployments (groups, aliases, wildcard patterns, team mappings) is the router registry, defined +//! in `litellm-router-core`, which reads this catalog through its dynamic config. +//! +//! # Structure +//! - `Catalog`: the parsed map. `Catalog::parse` takes the JSON body and a `Provenance` (where +//! the body came from: the bundled copy in the repo or the remote copy, with revision and etag), +//! and keeps every entry in file order with its raw fields. +//! - `ModelEntry`: one model. `ModelEntry::fields` is the raw JSON object, so unknown fields are +//! never lost; `ModelEntry::info` is the typed mirror `ModelInfo` (pricing in `pricing`, +//! modalities and modes in `capabilities`). +//! - Aliases: an entry may list `aliases`; each alias resolves to its canonical key. Problems +//! found while indexing them (bad lists, collisions) are kept as `AliasIssue`s instead of +//! failing the parse. +//! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known +//! families. They are data in the file, not entries. +//! - Validation: `Catalog::validate` checks a catalog against `IntegrityLimits` (a remote copy that +//! lost most of its models is refused), and `validate_registry` checks every entry against the +//! checked-in contract. +//! +//! # Lookups +//! `Catalog::lookup` resolves a key exactly first (entry or alias), then case-insensitively, and +//! returns a `ModelMatch` with both the key that matched and the canonical key of the entry. Provider +//! prefixes are not stripped here; callers decide which keys to try. +//! +//! # Readers +//! Cost tracking, the router (through `RouterDynamicConfig`) and the gateway. Nothing in this crate +//! performs I/O; loading and refreshing the file is the caller's job. + mod capabilities; mod catalog; mod error; diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml new file mode 100644 index 00000000000..af97bd52c57 --- /dev/null +++ b/litellm-rust/crates/router-core/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "litellm-router-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop" + +[dependencies] diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs new file mode 100644 index 00000000000..05d95d02c84 --- /dev/null +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -0,0 +1,63 @@ +//! Routing for LiteLLM: decide which deployment serves a request, and drive retries, fallbacks +//! and cooldowns around the attempts. +//! +//! This crate defines the boundary `litellm.Router` is moving behind. Today the Python router in +//! `litellm/router.py` implements all of it; this crate is where the Rust core will implement the +//! same interface. Until then it holds this overview and nothing else, so the interface is written +//! down in code before either side changes. +//! +//! # The public surface +//! Code outside the router uses `litellm.Router` through two kinds of members and nothing else: +//! - The endpoint methods (`acompletion`, `aembedding`, `aresponses`, `aanthropic_messages`, the +//! files, batches, assistants and vector store wrappers, 37 in all). Each is a thin wrapper over +//! `call(op, request)`. +//! - Eight entry points: `registry()`, `settings()`, `health()`, `select()`, `apply()`, +//! `observe()`, `register_plugin()` and `retire()`. +//! +//! Every other member that outside code uses today has a named replacement among those. The old +//! names stay available for about a month through a `DeprecatedRouterMembers` mixin, which +//! implements them over the entry points and warns on use. A CI check +//! (`tests/code_coverage_tests/router_surface_check.py`) keeps an allowlist of the members outside +//! code touches; the list may only shrink. +//! +//! # Structure +//! - `Registry`: the configured deployments, indexed by id, model group, alias, wildcard pattern +//! and team. Immutable and versioned. Seven queries: `deployment(id)`, `deployments(filter)`, +//! `names(filter)`, `resolve(name, scope)`, `group(name)`, `access_groups()` and +//! `credentials(id)`. "Registry" is the router's view of what is configured; per-model pricing +//! and capabilities are the model catalog (`litellm-model-catalog`), which the router reads +//! through its dynamic config. +//! - `RouterConfig` / `RouterDynamicConfig`: everything the router needs to know, passed in. The +//! router never reads process globals. Static config (model list, strategy, retry, timeout and +//! cooldown policy, Redis settings, plugins) is fixed for the router's life. Dynamic config +//! (pricing, model aliases, fallbacks, retry counts, callbacks) is a versioned snapshot read +//! once per call. +//! - `Settings`: one typed, frozen view of the configuration a caller may read. +//! - `Health`: cooldowns, blocked deployments and usage, kept in shared cache keys so Python and +//! Rust workers see the same state. `report(id, outcome)` records outcomes the router did not +//! see itself (a health check, an interceptor that marks a deployment failed). +//! - `select(model, request)`: resolve a name, filter eligible deployments, pick one with the +//! configured strategy, without calling it. +//! - `apply(changes)`: one `ChangeSet` (upserts, deletes, settings), applied atomically, producing +//! a new registry version. +//! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the +//! policy is exhausted. +//! +//! # Execution +//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs +//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A +//! started attempt is never replayed on a different executor. +//! +//! # Hosts +//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select` +//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks, +//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all +//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost` +//! serves them from the existing Python code when the router runs inside Python; `NativeHost` +//! serves them in the Rust server, where an operation without a Rust implementation makes the +//! config fail at build time, never partway through a request. +//! +//! # Plugins +//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive) +//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and +//! the current request), never the router itself. diff --git a/litellm/router_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py similarity index 100% rename from litellm/router_utils/batch_utils.py rename to litellm/litellm_core_utils/batch_utils.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index cf6bdc69556..101f069c0f8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -21,8 +21,8 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.batch_utils import InMemoryFile from litellm.litellm_core_utils.core_helpers import as_str_mapping -from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 7a4c707860d..17c178d169e 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.anthropic.messages.native.utils import AnthropicMessagesRequestUtils, prepare_native_messages -from litellm.router_utils.baseline_request import ( +from litellm.router_strategy.auto_router.baseline_request import ( BASELINE_PARAMETERS, capture_baseline_parameters, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 70e305f46fb..9bb451495c8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -139,8 +139,8 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ) from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict -from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 1267fb15db8..3071bcbb1dc 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket -from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 0d2f8dc9a0c..be32405154c 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -27,7 +27,7 @@ from litellm.constants import ( HEALTH_CHECK_TIMEOUT_SECONDS, ) from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( StrategyRouterDependency, classify_strategy_router_model, strategy_router_dependencies, diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py index c4232609124..1147ce7ec12 100644 --- a/litellm/proxy/hooks/autorouter_baseline_cache.py +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -37,8 +37,8 @@ from litellm.proxy.spend_tracking.savings import ( _effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner _proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner ) +from litellm.router_strategy.auto_router.baseline_request import baseline_request from litellm.router_strategy.complexity_router.context_compaction import compaction_applied -from litellm.router_utils.baseline_request import baseline_request from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import BaselineRouteStamp from litellm.types.utils import CallTypes, ModelInfo, Usage diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 11a47b2d424..1c4a61cfad8 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,6 +19,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span +from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions @@ -30,7 +31,6 @@ from litellm.proxy.hooks.model_budget_shared_windows import ( SharedWindows, ) from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds -from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d4528bc85e8..2c05f724374 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3290,7 +3290,7 @@ def _per_user_credential_names_for_groups( GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, ) - from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model + from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model names: Final[list[str]] = [] # mutable-ok: accumulates one name per per-user deployment for group in model_groups: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 4f48f3c57b9..3dceab6b7f7 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -60,16 +60,16 @@ from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + StrategyRouterDependencyRole, + classify_strategy_router_model, + strategy_router_dependencies, +) from litellm.router_strategy.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.request_models import ( AutoRouterRoutingTestRequest, RequestComplexityRouterConfig, ) -from litellm.router_utils.auto_router_model_naming import ( - StrategyRouterDependencyRole, - classify_strategy_router_model, - strategy_router_dependencies, -) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.management_endpoints.auto_router_endpoints import ( SHADOW_EVAL_TURN_VALVE, @@ -411,7 +411,7 @@ async def validate_complexity_router_config( auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict) ) - from litellm.router_utils.auto_router_model_naming import ( + from litellm.router_strategy.auto_router.auto_router_model_naming import ( validate_complexity_router_config_write, ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 45006b14df6..8cd3559ba86 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -131,6 +131,23 @@ from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + carries_complexity_router_settings, + count_capability_routers, + gated_capabilities_of, + is_complexity_router_model, + validate_complexity_router_config_placement, + validate_complexity_router_config_write, + validate_strategy_router_model_write, +) +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( + is_mutable_tuned_candidate, + tuning_quota_violation, +) from litellm.router_strategy.complexity_router import ( DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, ClassificationRubric, @@ -144,20 +161,6 @@ from litellm.router_strategy.complexity_router import ( normalize_classification_prompt, ) from litellm.router_strategy.complexity_router.config import resolve_complexity_router_config_write -from litellm.router_utils.auto_router_model_naming import ( - GATED_AUTO_ROUTER_CAPABILITIES, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - carries_complexity_router_settings, - count_capability_routers, - gated_capabilities_of, - is_complexity_router_model, - validate_complexity_router_config_placement, - validate_complexity_router_config_write, - validate_strategy_router_model_write, -) -from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsSessionTag from litellm.types.proxy.management_endpoints.model_management_endpoints import ( diff --git a/litellm/proxy/management_helpers/auto_router_availability.py b/litellm/proxy/management_helpers/auto_router_availability.py index 9b511ce0817..ce7859aff2d 100644 --- a/litellm/proxy/management_helpers/auto_router_availability.py +++ b/litellm/proxy/management_helpers/auto_router_availability.py @@ -7,14 +7,14 @@ from typing import Final from pydantic import Json, TypeAdapter, ValidationError from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, capability_limit_violation, classify_strategy_router_model, count_capability_routers, gated_capabilities_of, ) -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( is_mutable_tuned_candidate, mutable_tuned_identities, tuning_quota_violation, diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index 995446db7a6..8597f0d3667 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -34,13 +34,16 @@ from litellm.repositories.prisma_protocols import DatabaseClient from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + classify_strategy_router_model, + strategy_router_dependencies, +) from litellm.router_strategy.complexity_router.config import ( ComplexityRouterConfigWrite, normalize_classifier_config_aliases, resolve_complexity_router_config_write, ) from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig -from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model, strategy_router_dependencies from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import Deployment, updateDeployment diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 70e1c25a70a..b263412ce0f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -168,12 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields -from litellm.router_utils.access_windows import access_windows_config_error -from litellm.router_utils.add_retry_fallback_headers import ( - get_fallback_errors_from_headers, - get_hidden_params_dict, -) -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, capability_limit_violation, @@ -181,12 +176,17 @@ from litellm.router_utils.auto_router_model_naming import ( count_capability_routers, validate_complexity_router_config_placement, ) -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( TUNING_BASELINE_PARAM_NAME, mutable_tuned_identities, snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.router_utils.access_windows import access_windows_config_error +from litellm.router_utils.add_retry_fallback_headers import ( + get_fallback_errors_from_headers, + get_hidden_params_dict, +) from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation diff --git a/litellm/router.py b/litellm/router.py index 408cb20945f..fc888495e99 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7,6 +7,29 @@ # # Thank you ! We ❤️ you! - Krrish & Ishaan +""" +litellm.Router: load balancing, retries, fallbacks and cooldowns across deployments. + +Public surface. The router is moving behind a narrow interface; code outside +the router should end up using a Router through two kinds of members only: + +- endpoint methods (acompletion, aembedding, aresponses, aanthropic_messages, + the files, batches, assistants and vector store wrappers, ...), each a thin + wrapper around one call; +- entry points (being added) that replace direct access to internals: registry() for + deployment lookups, settings() for configuration reads, health() for + cooldowns, blocks and usage, select() to pick a deployment without calling + it, apply() for runtime changes, observe() for routing events, + register_plugin() for strategies and sub-routers, and retire() to release + the router. + +Everything else on the class is internal. The members outside code uses today +are listed in tests/code_coverage_tests/router_surface_allowlist.json, and a CI +check (router_surface_check.py next to it) fails on any new outside use, so the +list only shrinks while callers move to the entry points. The same interface +is documented for the Rust core in litellm-rust/crates/router-core. +""" + import asyncio import contextlib import copy @@ -85,6 +108,15 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.batch_utils import ( + batch_create_routing, + deployment_supported_endpoints, + get_router_metadata_variable_name, + is_batch_retrieve_call_type, + replace_model_in_jsonl, + should_replace_model_in_jsonl, + without_deployments_that_cannot_batch, +) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, @@ -148,6 +180,15 @@ from litellm.llms.openai_like.model_info import ( get_openai_compatible_model_info, ) from litellm.responses.bridge_check import keep_registered_mode +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + AUTO_ROUTER_MODEL_PREFIX, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capabilities, + classify_strategy_router_model, + count_capability_routers, +) from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.context_compaction import ( @@ -187,24 +228,6 @@ from litellm.router_utils.add_retry_fallback_headers import ( replace_complexity_router_headers, response_total_token_count, ) -from litellm.router_utils.auto_router_model_naming import ( - AUTO_ROUTER_MODEL_PREFIX, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - claimed_capabilities, - classify_strategy_router_model, - count_capability_routers, -) -from litellm.router_utils.batch_utils import ( - batch_create_routing, - deployment_supported_endpoints, - get_router_metadata_variable_name, - is_batch_retrieve_call_type, - replace_model_in_jsonl, - should_replace_model_in_jsonl, - without_deployments_that_cannot_batch, -) from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit from litellm.router_utils.clientside_credential_handler import ( get_dynamic_litellm_params, @@ -16271,7 +16294,7 @@ class Router: to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ - from litellm.router_utils.baseline_request import capture_baseline_parameters + from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters from litellm.types.router import BaselineRouteStamp phase_attributes(routing_decision_attributes(routing_decision)) diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_strategy/auto_router/auto_router_model_naming.py similarity index 100% rename from litellm/router_utils/auto_router_model_naming.py rename to litellm/router_strategy/auto_router/auto_router_model_naming.py diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_strategy/auto_router/auto_router_tuning_baseline.py similarity index 100% rename from litellm/router_utils/auto_router_tuning_baseline.py rename to litellm/router_strategy/auto_router/auto_router_tuning_baseline.py diff --git a/litellm/router_utils/baseline_request.py b/litellm/router_strategy/auto_router/baseline_request.py similarity index 100% rename from litellm/router_utils/baseline_request.py rename to litellm/router_strategy/auto_router/baseline_request.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index eab232eae18..1d502dd8d95 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,12 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs -from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 6f3e0936641..8e0c3d6a7e7 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 5894da9cb4f..17d03fd4ff8 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -10,7 +10,7 @@ from litellm._internal_context import with_service_target 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.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 4e994d3fdc6..c352934db6c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -13,9 +13,9 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import LiteLLMPydanticObjectBase diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 2d373e0c266..9636d7f009d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -9,7 +9,7 @@ from litellm._internal_context import with_service_target 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.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 527ea830d55..0719b30d36d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -15,8 +15,8 @@ from litellm._internal_context import with_service_target 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.batch_utils import is_batch_retrieve_call_type 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 diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0a77ad5b9ea..6d1de1f9f6c 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -13,6 +13,7 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider @@ -22,7 +23,6 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, get_fallback_error_info, ) -from litellm.router_utils.batch_utils import get_router_metadata_variable_name from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 4c9afd26dad..9fb3633d53c 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -187,13 +187,13 @@ ignored_function_names = [ def main(): router_file = [ "./litellm/router.py", - "./litellm/router_utils/batch_utils.py", + "./litellm/litellm_core_utils/batch_utils.py", "./litellm/router_utils/pattern_match_deployments.py", ] # router_file = [ # "../../litellm/router.py", # "../../litellm/router_utils/pattern_match_deployments.py", - # "../../litellm/router_utils/batch_utils.py", + # "../../litellm/litellm_core_utils/batch_utils.py", # ] ## LOCAL TESTING tests_dir = "./tests/" # Update this path if your tests directory is located elsewhere # tests_dir = "../../tests/" # LOCAL TESTING diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json new file mode 100644 index 00000000000..428fd9b4011 --- /dev/null +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -0,0 +1,211 @@ +{ + "_doc": "Router members and litellm.router_utils modules used outside the router. Shrink-only: see tests/code_coverage_tests/router_surface_check.py.", + "members_outside": [ + "_get_client", + "_zero_cost_cache", + "a_add_message", + "aadapter_completion", + "aanthropic_messages", + "abatch_completion", + "abatch_completion_fastest_response", + "acancel_batch", + "acancel_fine_tuning_job", + "acompletion", + "acreate_assistants", + "acreate_batch", + "acreate_file", + "acreate_fine_tuning_job", + "acreate_thread", + "adaptive_routers", + "adecisions", + "adelete_assistant", + "aembedding", + "afile_content", + "afile_delete", + "aget_assistants", + "aget_messages", + "aget_responses", + "aget_thread", + "alist_batches", + "alist_fine_tuning_jobs", + "allm_passthrough_route", + "amoderation", + "aocr", + "are_all_deployments_blocked", + "arefresh_model_info", + "aretrieve_batch", + "aretrieve_fine_tuning_job", + "arm_routing_read_prefetch", + "arun_thread", + "async_get_available_deployment", + "async_get_fully_unhealthy_model_names", + "async_get_healthy_deployments", + "async_pre_routing_hook", + "auto_routers", + "avector_store_create", + "avector_store_search", + "background_health_check_model_groups", + "bind_to_request_logging", + "cache", + "cache_responses", + "cached_deployment_model_info", + "cached_model_group_info", + "completion", + "complexity_routers", + "config_deployments", + "content_policy_fallbacks", + "context_window_fallbacks", + "cooldown_time", + "default_deployment", + "default_litellm_params", + "delete_deployment", + "deployment_latency_map", + "deployment_names", + "deployment_usable_by_team", + "deployments_for_request", + "deployments_named", + "discard", + "embedding", + "enable_health_check_routing", + "fallbacks", + "generate_model_id", + "get_all_deployments", + "get_available_deployment_for_pass_through", + "get_available_guardrail", + "get_candidate_model_ids_for_route", + "get_configured_display_name", + "get_configured_mode", + "get_configured_service_tiers", + "get_configured_token_limits", + "get_credential_deployment", + "get_deployment", + "get_deployment_by_model_group_name", + "get_deployment_credentials", + "get_deployment_credentials_with_provider", + "get_deployment_model_info", + "get_discovered_model_info", + "get_fully_blocked_model_names", + "get_model_access_groups", + "get_model_from_alias", + "get_model_group_info", + "get_model_group_usage", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_model_list_from_model_alias", + "get_model_list_from_routing_groups", + "get_model_list_of_routed_group", + "get_model_listing_info", + "get_model_names", + "get_remaining_model_group_usage", + "get_routable_deployments", + "get_routable_upstream_model", + "get_router_deployment_budget_limiter", + "get_routing_group", + "get_settings", + "get_valid_args", + "guardrail_list", + "has_model_id", + "health_check_ignore_transient_errors", + "health_state_cache", + "is_deployment_blocked", + "is_model_access_group_for_wildcard_route", + "is_recognized_model", + "kubernetes_pod_discovery", + "map_team_model", + "model_group_alias", + "model_list", + "model_name_to_deployment_indices", + "model_names", + "pattern_router", + "priority_header_update", + "prompt_caching_pin_cache", + "provider_budget_config", + "quality_routers", + "resolve_model_name_from_model_id", + "resolve_served_model_name", + "resolved_litellm_models", + "routable_model_group", + "router_general_settings", + "routing_groups", + "routing_plugins", + "schedule_acompletion", + "search_tools", + "serve_through_request_logging", + "should_include_deployment", + "split_request_priority", + "team_model_to_deployment_indices", + "team_pattern_routers", + "team_ptu_ceiling", + "team_public_model_names", + "team_public_model_names_of", + "update_redis_cache", + "update_settings", + "upsert_deployment", + "wildcard_route_deployments", + "wildcard_routed_models" + ], + "members_router_family": [ + "_count_pre_call_check_tokens", + "_pre_call_checks", + "_update_kwargs_before_fallbacks", + "acompletion", + "adecisions", + "aembedding", + "allowed_fails", + "allowed_fails_policy", + "aresponses", + "async_function_with_fallbacks", + "async_get_healthy_deployments", + "cache", + "cached_deployment_model_info", + "cooldown_cache", + "cooldown_time", + "default_max_parallel_requests", + "deployments_named", + "disable_cooldowns", + "embedding", + "enable_pre_call_checks", + "enable_tag_filtering", + "explicit_allowed_fails", + "fallback_access_check", + "fallback_budget_check", + "fallbacks", + "get_all_deployments", + "get_allowed_fails_from_policy", + "get_deployment", + "get_model_group", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_router_model_info", + "log_retry", + "model_list", + "provider_default_deployment_ids", + "routing_group_has_alternatives", + "routing_strategy", + "search_tools", + "set_custom_routing_strategy", + "slack_alerting_logger", + "tag_filtering_match_any", + "tag_routing_prefix", + "team_model_has_alternatives" + ], + "router_utils_imports_outside": [ + "litellm.router_utils.access_windows", + "litellm.router_utils.add_retry_fallback_headers", + "litellm.router_utils.clientside_credential_handler", + "litellm.router_utils.common_utils", + "litellm.router_utils.cooldown_handlers", + "litellm.router_utils.deployment_pricing", + "litellm.router_utils.fallback_event_handlers", + "litellm.router_utils.get_retry_from_policy", + "litellm.router_utils.health_state_cache", + "litellm.router_utils.pattern_match_deployments", + "litellm.router_utils.ptu_shares", + "litellm.router_utils.reasoning_effort_capability", + "litellm.router_utils.router_callbacks.track_deployment_metrics", + "litellm.router_utils.routing_groups", + "litellm.router_utils.search_api_router" + ] +} diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py new file mode 100644 index 00000000000..0b35bac5240 --- /dev/null +++ b/tests/code_coverage_tests/router_surface_check.py @@ -0,0 +1,235 @@ +""" +Fence for the Router's public surface. + +Code outside the router may use only the Router members listed in +router_surface_allowlist.json, and may import only the litellm.router_utils +modules listed there. The lists hold what is used today and may only shrink: + +- a member or module used outside the router that is not listed fails the + check (use an existing entry point instead, or discuss widening the surface); +- a listed member or module that nothing outside the router uses any more + also fails, so the entry is removed in the same change. + +"Outside" means every non-test file under litellm/ and enterprise/ except +litellm/router.py and the router's own modules. Uses from the router family +(litellm/router_utils, litellm/router_strategy) are listed separately, because +those modules become private parts of the router core in a later step. + +Member uses are found the same way the boundary survey found them: attribute +access on a receiver whose name looks like a router (llm_router, +litellm_router_instance, router, ...) or on the Router class, plus +getattr/setattr/hasattr with a constant name on such a receiver. Members +include the container file endpoints Router.__init__ registers by name. + +Run: python tests/code_coverage_tests/router_surface_check.py [--write] +--write regenerates the allowlist from the current tree (for shrinking it). +""" + +import ast +import json +import os +import re +import sys +from typing import Dict, List, Set + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.json") + +ROUTER_FILES = { + "litellm/router.py", +} +ROUTER_CLASS_FILES = ["litellm/router.py"] +ROUTER_CLASS_NAMES = {"Router"} + +SCAN_ROOTS = ["litellm", "enterprise"] +SKIP_DIRS = { + ".git", + "node_modules", + ".venv", + "__pycache__", + "litellm-rust", + "ui", + "docs", + "cookbook", + "tests", + "test", +} + +RECEIVER = re.compile( + r"(^|[._])(llm_router|router|litellm_router|litellm_router_instance|router_instance" + r"|_router|proxy_router|llm_router_instance|router_obj|mock_router|Router)$" +) +# Router.__init__ registers these with setattr(self, name, ...), so they never appear as self.. +CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json" +DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"} + + +def _area(rel: str) -> str: + if rel.startswith(("litellm/router_utils/", "litellm/router_strategy/")): + return "router_family" + return "outside" + + +def _is_test(rel: str) -> bool: + name = os.path.basename(rel) + return "/tests/" in rel or name.startswith("test_") or name.endswith("_test.py") + + +def router_members() -> Set[str]: + members: Set[str] = set() + for rel in ROUTER_CLASS_FILES: + with open(os.path.join(ROOT, rel), encoding="utf-8") as f: + tree = ast.parse(f.read()) + for cls in tree.body: + if not isinstance(cls, ast.ClassDef) or cls.name not in ROUTER_CLASS_NAMES: + continue + for node in cls.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + members.add(node.name) + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + members.add(node.target.id) + elif isinstance(node, ast.Assign): + members.update(t.id for t in node.targets if isinstance(t, ast.Name)) + for node in ast.walk(cls): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + and isinstance(node.ctx, ast.Store) + ): + members.add(node.attr) + with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f: + for endpoint in json.load(f)["endpoints"]: + members.update((endpoint["name"], endpoint["async_name"])) + return {m for m in members if not (m.startswith("__") and m.endswith("__"))} + + +def _python_files(): + for top in SCAN_ROOTS: + base = os.path.join(ROOT, top) + for dirpath, dirnames, filenames in os.walk(base): + dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS] + for fn in filenames: + if fn.endswith(".py"): + path = os.path.join(dirpath, fn) + yield os.path.relpath(path, ROOT).replace(os.sep, "/"), path + + +def scan() -> Dict[str, Dict[str, Set[str]]]: + """Return {"members": {area: {member: files}}, "router_utils_imports": {module: files}}.""" + members = router_members() + uses: Dict[str, Dict[str, Set[str]]] = {"outside": {}, "router_family": {}} + imports: Dict[str, Set[str]] = {} + for rel, path in _python_files(): + if rel in ROUTER_FILES or _is_test(rel): + continue + with open(path, encoding="utf-8") as f: + src = f.read() + has_router = "router" in src.lower() + if not has_router: + continue + try: + tree = ast.parse(src) + except SyntaxError: + continue + area = _area(rel) + for name in member_uses(tree, members): + uses[area].setdefault(name, set()).add(rel) + if area == "outside": + for node in ast.walk(tree): + for mod in _router_utils_modules(node): + imports.setdefault(mod, set()).add(rel) + return {"members": uses, "router_utils_imports": imports} + + +def member_uses(tree: ast.AST, members: Set[str]) -> Set[str]: + """The Router members a module uses, on an instance or on the Router class itself.""" + used: Set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and node.attr in members: + if RECEIVER.search(ast.unparse(node.value)): + used.add(node.attr) + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id in DYNAMIC_FUNCS + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and node.args[1].value in members + and RECEIVER.search(ast.unparse(node.args[0])) + ): + used.add(node.args[1].value) + return used + + +def _router_utils_modules(node: ast.AST) -> List[str]: + if isinstance(node, ast.ImportFrom) and node.module: + if node.module == "litellm.router_utils": + return ["litellm.router_utils." + a.name for a in node.names] + if node.module.startswith("litellm.router_utils."): + return [node.module] + if isinstance(node, ast.Import): + return [a.name for a in node.names if a.name.startswith("litellm.router_utils.")] + return [] + + +def build_allowlist(found) -> dict: + return { + "_doc": ( + "Router members and litellm.router_utils modules used outside the router. " + "Shrink-only: see tests/code_coverage_tests/router_surface_check.py." + ), + "members_outside": sorted(found["members"]["outside"]), + "members_router_family": sorted(found["members"]["router_family"]), + "router_utils_imports_outside": sorted(found["router_utils_imports"]), + } + + +def compare(allow: dict, found) -> List[str]: + errors: List[str] = [] + sections = [ + ("members_outside", found["members"]["outside"], "Router member"), + ("members_router_family", found["members"]["router_family"], "Router member (router family)"), + ("router_utils_imports_outside", found["router_utils_imports"], "litellm.router_utils import"), + ] + for key, used, what in sections: + listed = set(allow.get(key, [])) + for name in sorted(set(used) - listed): + files = ", ".join(sorted(used[name])) + errors.append( + f"new {what} outside the router: {name!r} in {files}. The router surface may only " + f"shrink: use an entry point instead." + ) + for name in sorted(listed - set(used)): + errors.append( + f"{what} {name!r} is listed in {os.path.basename(ALLOWLIST)} [{key}] but nothing uses it " + f"any more: remove it from the list." + ) + return errors + + +def main(argv: List[str]) -> int: + found = scan() + if "--write" in argv: + with open(ALLOWLIST, "w", encoding="utf-8") as f: + json.dump(build_allowlist(found), f, indent=2) + f.write("\n") + print(f"wrote {ALLOWLIST}") + return 0 + with open(ALLOWLIST, encoding="utf-8") as f: + allow = json.load(f) + errors = compare(allow, found) + for e in errors: + print(e) + if errors: + return 1 + print( + f"router surface ok: {len(allow['members_outside'])} members outside, " + f"{len(allow['members_router_family'])} from the router family, " + f"{len(allow['router_utils_imports_outside'])} router_utils modules imported outside" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt index 1da964681a9..d8c6e2241f4 100644 --- a/tests/code_coverage_tests/unbounded_in_baseline.txt +++ b/tests/code_coverage_tests/unbounded_in_baseline.txt @@ -129,4 +129,4 @@ litellm/proxy/utils.py prefetch_config_params prisma param_name.in `param_names` litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma [scope.entity_id_field].in `list(scope.entity_ids)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma api_key.in `list(scope.api_keys)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma not.in `list(scope.exclude_entity_ids)` 0 -litellm/router_utils/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 +litellm/router_strategy/auto_router/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 diff --git a/tests/unit/router_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py similarity index 97% rename from tests/unit/router_utils/test_batch_utils.py rename to tests/unit/litellm_core_utils/test_batch_utils.py index 79fe86a37c0..cd37d635c43 100644 --- a/tests/unit/router_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -7,7 +7,7 @@ from typing import Dict, Final, List import pytest import litellm -from litellm.router_utils.batch_utils import ( +from litellm.litellm_core_utils.batch_utils import ( InMemoryFile, get_router_metadata_variable_name, parse_jsonl_with_embedded_newlines, @@ -116,7 +116,7 @@ def test_router_metadata_variable_name(): def test_non_json_input(): """Test that replace_model_in_jsonl returns original content for non-JSON input""" - from litellm.router_utils.batch_utils import replace_model_in_jsonl + from litellm.litellm_core_utils.batch_utils import replace_model_in_jsonl # Test with non-JSON string non_json_str = "This is not a JSON string" @@ -138,7 +138,7 @@ def test_non_json_input(): def test_should_replace_model_in_jsonl(): """Test that should_replace_model_in_jsonl returns the correct value""" - from litellm.router_utils.batch_utils import should_replace_model_in_jsonl + from litellm.litellm_core_utils.batch_utils import should_replace_model_in_jsonl assert should_replace_model_in_jsonl(purpose="batch") is True assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False @@ -340,7 +340,7 @@ def test_replace_model_in_jsonl_with_embedded_newlines(): 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.litellm_core_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 @@ -389,7 +389,8 @@ def setup_and_teardown(): ) def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: from openai.types import Batch - from litellm.router_utils.batch_utils import batch_id_to_charge_once + + from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once batch: Final = Batch( id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", @@ -401,7 +402,7 @@ def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache - from litellm.router_utils.batch_utils import batch_charged_once_marker_key + from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key cache: Final = InMemoryCache() cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py index 748b291d5ff..c13e3d5812b 100644 --- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py @@ -170,7 +170,7 @@ async def test_preview_and_serving_share_default_model_resolution( explicit_default: str | None, expected: str, ): - from litellm.router_utils.auto_router_model_naming import validate_complexity_router_config_write + from litellm.router_strategy.auto_router.auto_router_model_naming import validate_complexity_router_config_write from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest config: Final = { @@ -1240,7 +1240,7 @@ class TestAutoRouterBenchmarks: def test_the_listed_kinds_match_the_router_types_traffic_can_record(self): from typing import get_args, get_type_hints - from litellm.router_utils.auto_router_model_naming import StrategyRouterKind + from litellm.router_strategy.auto_router.auto_router_model_naming import StrategyRouterKind from litellm.types.utils import StandardLoggingRoutingDecision readonly_router_type: Final = get_type_hints(StandardLoggingRoutingDecision, include_extras=True)["router_type"] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index fc6749aed0f..437d948f8f5 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1,6 +1,6 @@ -import inspect import asyncio import contextlib +import inspect import json import socket from collections.abc import Iterator, Mapping @@ -1461,10 +1461,10 @@ class TestDeleteModelClearsRouterRegistry: add_deployment never restores, so an unguarded pop would make it permanently unroutable (the same cross-tenant DoS clear_cache was hardened against). """ + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) - from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete model_id = "regular-del-1" admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) @@ -2739,11 +2739,11 @@ class TestAddAndDeleteModelLifecycle: - Delete same model again → raises (model not found) """ from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, add_new_model, - delete_model as delete_model_endpoint, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( - ModelInfoDelete, + delete_model as delete_model_endpoint, ) model_id = "lifecycle-test-model-123" @@ -2833,6 +2833,8 @@ class TestDeleteTeamBYOKModelGhost: async def test_delete_strips_public_name_and_refreshes_cache(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -2916,6 +2918,8 @@ class TestDeleteTeamBYOKModelGhost: alias cleanup (delete_team_model_alias), preserving legacy behavior.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -2984,6 +2988,8 @@ class TestDeleteTeamBYOKModelGhost: team.models when one replica is deleted but a sibling still backs it.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3065,6 +3071,8 @@ class TestDeleteTeamBYOKModelGhost: deployment still serves it.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3145,6 +3153,8 @@ class TestDeleteTeamBYOKModelGhost: the alias would break routing that works.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3252,6 +3262,8 @@ class TestDeleteModelTeamAuth: async def test_proxy_admin_can_delete_model_when_team_deleted(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3287,6 +3299,8 @@ class TestDeleteModelTeamAuth: """A missing team must never let a non-admin delete the orphan (no fail-open).""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) from litellm.proxy.proxy_server import ProxyException @@ -3322,6 +3336,8 @@ class TestDeleteModelTeamAuth: """The auth check must not add a redundant team query on the live-team path.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) from litellm.proxy.proxy_server import ProxyException @@ -4392,7 +4408,6 @@ class TestModelInfoCostMapEchoFilter: def test_echoed_cost_map_metadata_is_not_persisted(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4418,7 +4433,6 @@ class TestModelInfoCostMapEchoFilter: def test_an_edited_value_survives_the_echo_filter(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4483,7 +4497,6 @@ class TestModelInfoCostMapEchoFilter: def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import echoed_cost_map_fields @@ -4513,7 +4526,6 @@ class TestModelInfoCostMapEchoFilter: def test_a_stored_mode_survives_an_echoed_save(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4536,7 +4548,6 @@ class TestModelInfoCostMapEchoFilter: def test_resetting_an_override_to_the_cost_map_value_removes_it(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4560,7 +4571,6 @@ class TestModelInfoCostMapEchoFilter: def test_reset_is_recognised_after_the_router_registered_the_override(self, monkeypatch: pytest.MonkeyPatch): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4586,7 +4596,6 @@ class TestModelInfoCostMapEchoFilter: from types import MappingProxyType import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4611,7 +4620,6 @@ class TestModelInfoCostMapEchoFilter: def test_echo_is_compared_against_the_deployments_lookup_not_the_key(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4653,7 +4661,6 @@ class TestModelInfoCostMapEchoFilter: def test_base_model_wins_over_litellm_params_model_for_the_lookup(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4681,7 +4688,6 @@ class TestModelInfoCostMapEchoFilter: def test_encrypted_stored_model_is_decrypted_for_the_lookup(self, monkeypatch): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -6431,7 +6437,7 @@ class TestStrategyRouterWriteValidation: AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY, _auto_router_capability_slot, ) - from litellm.router_utils.auto_router_model_naming import gated_capabilities_of + from litellm.router_strategy.auto_router.auto_router_model_naming import gated_capabilities_of capabilities: Final = gated_capabilities_of(effective_params) @@ -6568,7 +6574,7 @@ class TestStrategyRouterWriteValidation: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines configs = { "_TUNED_A": self._TUNED_A, @@ -6624,7 +6630,7 @@ class TestStrategyRouterWriteValidation: from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines fake = self._FakeDb([]) fake.tx_obj.litellm_proxymodeltable.find_many = AsyncMock( @@ -6669,7 +6675,7 @@ class TestStrategyRouterWriteValidation: row id before the slot is entered, so a second tuned router is refused before its DB write.""" from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)]) @@ -8980,7 +8986,10 @@ class TestTeamMemberAutoRouterWrites: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.proxy.management_helpers.auto_router_permissions import MemberAutoRouterWrite, validate_member_auto_router_config + from litellm.proxy.management_helpers.auto_router_permissions import ( + MemberAutoRouterWrite, + validate_member_auto_router_config, + ) row: Final = self._row() database: Final = self._database(self._team(), row) @@ -10677,6 +10686,8 @@ class TestFederationGateScopesToWhatTheWriteTouches: and the team that owns the model must be able to take it off their page.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) diff --git a/tests/unit/proxy/management_helpers/test_auto_router_availability.py b/tests/unit/proxy/management_helpers/test_auto_router_availability.py index a88dcd7f66d..0eaae60b4ea 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_availability.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_availability.py @@ -1,16 +1,15 @@ from collections.abc import Mapping -from typing import Final from types import SimpleNamespace - -from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from typing import Final import pytest +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.management_helpers.auto_router_availability import ( auto_router_availability, build_auto_router_catalog, ) -from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines def deployment( diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..8b1d7c1222a 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -33,9 +33,9 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( ProxyStartupEvent, - initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, cleanup_router_config_variables, @@ -43,6 +43,7 @@ from litellm.proxy.proxy_server import ( get_litellm_model_info, initialize, initialize_from_worker_config, + initialize_shared_aiohttp_session, load_from_azure_key_vault, proxy_shutdown_event, proxy_startup_event, @@ -50,8 +51,6 @@ from litellm.proxy.proxy_server import ( save_worker_config, ) -from litellm.proxy.management_helpers.audit_logs import track_audit_task - from .conftest import normalize # --------------------------------------------------------------------------- @@ -1165,7 +1164,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): @pytest.mark.asyncio async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row(): - from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT prisma_client = MagicMock() prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) @@ -1186,7 +1185,10 @@ async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row(): @pytest.mark.asyncio async def test_scorer_baseline_upgrade_preserves_existing_routers_and_is_not_refreshed_on_restart(): - from litellm.router_utils.auto_router_tuning_baseline import mutable_tuned_identities, snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( + mutable_tuned_identities, + snapshot_tuning_baselines, + ) deployments = [ { diff --git a/tests/unit/router_utils/test_auto_router_model_naming.py b/tests/unit/router_strategy/test_auto_router_model_naming.py similarity index 99% rename from tests/unit/router_utils/test_auto_router_model_naming.py rename to tests/unit/router_strategy/test_auto_router_model_naming.py index 1cfa3aa5c91..44fcfdaf5be 100644 --- a/tests/unit/router_utils/test_auto_router_model_naming.py +++ b/tests/unit/router_strategy/test_auto_router_model_naming.py @@ -3,9 +3,7 @@ from typing import Final import pytest -from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets -from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, capability_limit_violation, carries_complexity_router_settings, @@ -18,6 +16,8 @@ from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, validate_strategy_router_model_write, ) +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets +from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS COMPLEXITY_FIELDS = frozenset({"complexity_router_config"}) SEMANTIC_FIELDS = frozenset({"auto_router_config", "auto_router_default_model", "auto_router_embedding_model"}) @@ -596,7 +596,7 @@ def test_custom_classifier_prompt_capability(config: Mapping[str, object], expec ], ) def test_is_complexity_router_model(model: str | None, expected: bool) -> None: - from litellm.router_utils.auto_router_model_naming import is_complexity_router_model + from litellm.router_strategy.auto_router.auto_router_model_naming import is_complexity_router_model assert is_complexity_router_model(model) is expected diff --git a/tests/unit/router_utils/test_auto_router_tuning_baseline.py b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py similarity index 99% rename from tests/unit/router_utils/test_auto_router_tuning_baseline.py rename to tests/unit/router_strategy/test_auto_router_tuning_baseline.py index d0d98ce7042..4ee219a7784 100644 --- a/tests/unit/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py @@ -7,8 +7,7 @@ from typing import Final import pytest -from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( DEFAULT_TUNING_FINGERPRINT, HEURISTIC_V1_TUNING_FIELDS, heuristic_v1_router_fingerprint, @@ -19,6 +18,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( tuning_limit_violation, tuning_quota_violation, ) +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig _TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} _ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} diff --git a/tests/unit/router_utils/test_baseline_request.py b/tests/unit/router_strategy/test_baseline_request.py similarity index 95% rename from tests/unit/router_utils/test_baseline_request.py rename to tests/unit/router_strategy/test_baseline_request.py index 37d21597752..eb785a41d63 100644 --- a/tests/unit/router_utils/test_baseline_request.py +++ b/tests/unit/router_strategy/test_baseline_request.py @@ -1,6 +1,6 @@ from typing import Final -from litellm.router_utils.baseline_request import baseline_request, capture_baseline_parameters +from litellm.router_strategy.auto_router.baseline_request import baseline_request, capture_baseline_parameters def test_baseline_snapshot_owns_nested_caller_settings_and_overrides_routed_settings() -> None: diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ca39c7f5c8f..018f93779a6 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -24,14 +24,6 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm import Router -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingRoutingDecision -from litellm.router_utils.auto_router_model_naming import ( - CUSTOMIZATION_CAPABILITY, - GATED_AUTO_ROUTER_CAPABILITIES, - HEURISTIC_V2_CAPABILITY, - count_capability_routers, -) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache @@ -42,7 +34,19 @@ from litellm.constants import ( RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + CUSTOMIZATION_CAPABILITY, + GATED_AUTO_ROUTER_CAPABILITIES, + HEURISTIC_V2_CAPABILITY, + count_capability_routers, +) +from litellm.router_strategy.complexity_router.capability_classifier import ( + CAPABILITY_CLASSIFIER_SYSTEM_PROMPT, + CapabilityClassifierVerdict, +) from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -58,18 +62,14 @@ from litellm.router_strategy.complexity_router.complexity_router import ( classification_system_prompt, custom_tier_classification_prompt, ) -from litellm.router_strategy.complexity_router.capability_classifier import ( - CAPABILITY_CLASSIFIER_SYSTEM_PROMPT, - CapabilityClassifierVerdict, -) from litellm.router_strategy.complexity_router.config import ( - CapabilityCalibrationConfig, - CapabilityClassifierConfig, DEFAULT_CLASSIFICATION_RUBRIC, DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_COMPLEXITY_CONFIG, DEFAULT_TECHNICAL_KEYWORDS, TIER_SEVERITY_ORDER, + CapabilityCalibrationConfig, + CapabilityClassifierConfig, ClassificationRubric, ClassifierLLMConfig, ComplexityRouterConfig, @@ -83,10 +83,12 @@ from litellm.router_strategy.complexity_router.jev_classifier import ( JevUsage, ) from litellm.router_strategy.complexity_router.llm_v2 import LLM_V2_PROMPT_VERSION +from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, TrainedTierArtifact, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import ( Deployment, LiteLLM_Params, @@ -94,10 +96,7 @@ from litellm.types.router import ( RouterErrors, TaggedPreRoutingStrategy, ) -from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - +from litellm.types.utils import StandardLoggingRoutingDecision requires_semantic_router = pytest.mark.skipif( sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" @@ -16299,7 +16298,6 @@ class TestTierHealthFailover: tests can prove the probe passes a kwargs copy and forwards the prompt arguments. """ import litellm as litellm_module - from litellm.types.router import RouterRateLimitError probed_kwargs = [] diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py index 3fd2e8808e7..8c1e9e48661 100644 --- a/tests/unit/router_strategy/test_llm_v2.py +++ b/tests/unit/router_strategy/test_llm_v2.py @@ -4,11 +4,12 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -import litellm from pydantic import ValidationError +import litellm from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache +from litellm.router_strategy.auto_router.auto_router_model_naming import strategy_router_dependencies 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 @@ -21,7 +22,6 @@ from litellm.router_strategy.complexity_router.llm_v2 import ( LLMV2Verdict, llm_v2_response_format, ) -from litellm.router_utils.auto_router_model_naming import strategy_router_dependencies from litellm.types.llms.openai import ResponsesAPIResponse diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index 1d9c2e87100..d0c5c8247d3 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1969,7 +1969,7 @@ ALIAS_CASES: Final = ( ("litellm.router_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True), ("litellm.router_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False), ( - "litellm.router_utils.batch_utils", + "litellm.litellm_core_utils.batch_utils", "", "_get_router_metadata_variable_name", "get_router_metadata_variable_name", diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index af8e7f53462..513136ae150 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,20 +4,21 @@ import gc import json import os import weakref -from datetime import datetime, timezone from collections.abc import Iterator, Mapping +from datetime import datetime, timezone from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -import litellm import pytest + +import litellm from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY -from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.types.litellm_params import RoutingStrategyName @@ -801,10 +802,11 @@ async def test_deployment_callback_on_failure(model_list): def test_deployment_callback_respects_cooldown_time(model_list): """Ensure per-model cooldown_time is honored even when exception headers are present.""" - import httpx import time from unittest.mock import patch + import httpx + router = Router(model_list=model_list) class FakeException(Exception): @@ -1811,8 +1813,8 @@ def test_get_allowed_fails_from_policy(model_list, exception_type, exception_nam def test_initialize_alerting(model_list): """Test if the 'initialize_alerting' function is working correctly""" - from litellm.types.router import AlertingConfig from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.types.router import AlertingConfig router = Router( model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") @@ -2245,6 +2247,7 @@ def test_deployments_by_pattern(model_list): ) def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): from re import Match + from litellm.router_utils.pattern_match_deployments import PatternMatchRouter pattern_router = PatternMatchRouter() @@ -2702,7 +2705,7 @@ def test_handle_clientside_credential_metadata_loading(model_list, function_name ) def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key): """Test that _handle_clientside_credential uses the correct metadata variable name based on function name""" - from litellm.router_utils.batch_utils import get_router_metadata_variable_name + from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name router = Router(model_list=model_list) diff --git a/tests/unit/test_router/test_router_surface_check.py b/tests/unit/test_router/test_router_surface_check.py new file mode 100644 index 00000000000..a97fd029380 --- /dev/null +++ b/tests/unit/test_router/test_router_surface_check.py @@ -0,0 +1,58 @@ +""" +The router surface fence (tests/code_coverage_tests/router_surface_check.py) must see every way +outside code reaches a Router member, or new dependencies slip past the shrink-only allowlist. +""" + +import ast +import importlib.util +import pathlib + +import pytest + +_CHECK = pathlib.Path(__file__).parents[2] / "code_coverage_tests" / "router_surface_check.py" + + +@pytest.fixture(scope="module") +def surface_check(): + spec = importlib.util.spec_from_file_location("router_surface_check", _CHECK) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _uses(surface_check, src: str) -> set: + return surface_check.member_uses(ast.parse(src), surface_check.router_members()) + + +def test_instance_and_dynamic_access_are_uses(surface_check): + assert _uses(surface_check, "llm_router.get_model_list()\ngetattr(router, 'model_list')") == { + "get_model_list", + "model_list", + } + + +def test_class_level_calls_are_uses(surface_check): + assert _uses(surface_check, "litellm.Router.get_valid_args()") == {"get_valid_args"} + assert _uses(surface_check, "from litellm import Router\nRouter.get_valid_args()") == {"get_valid_args"} + + +def test_container_file_endpoints_registered_by_name_are_members(surface_check): + members = surface_check.router_members() + assert {"aupload_container_file", "upload_container_file", "alist_container_files"} <= members + assert _uses(surface_check, "await llm_router.aupload_container_file(file=f)") == {"aupload_container_file"} + + +def test_unlisted_class_level_use_fails_the_fence(surface_check): + found = { + "members": {"outside": {"get_valid_args": {"litellm/new_module.py"}}, "router_family": {}}, + "router_utils_imports": {}, + } + allow = {"members_outside": [], "members_router_family": [], "router_utils_imports_outside": []} + errors = surface_check.compare(allow, found) + assert len(errors) == 1 + assert "'get_valid_args' in litellm/new_module.py" in errors[0] + + +def test_unrelated_receivers_are_not_uses(surface_check): + assert _uses(surface_check, "client.get_model_list()\nRouterConfig.model_list") == set() From 49ccce7813725cc87b64d1b0f1b097dca5c681af Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:12:07 -0700 Subject: [PATCH 4/9] test(extraction): audit #47473 (#48898) * test(extraction): audit #47473 * test(extraction): check over-budget extraction never reaches Cloudflare --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../management/test_model_health_check.py | 37 ++ ...test_cloudflare_clairic_extraction_wire.py | 418 ++++++++++++++++++ 2 files changed, 455 insertions(+) create mode 100644 tests/integration/providers/test_cloudflare_clairic_extraction_wire.py diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 57a05daca94..992c1e8e875 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -26,6 +26,20 @@ _DECISIONS_PROBE_REPLY: Final = JsonResponse( "usage": {"input_tokens": 10, "output_tokens": 1}, }, ) +_EXTRACTION_PROBE_REPLY: Final = JsonResponse( + content_type="application/json", + body={ + "result": { + "model": "clairic", + "answers": {"reachable": {"type": "text", "value": "health", "confidence": 0.9}}, + "truncated": False, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, +) _CONFIGURED_PROBE_REPLY: Final = JsonResponse( content_type="application/json", body={ @@ -287,6 +301,29 @@ def test_decisions_mode_health_check_resolves_the_mode_from_the_cost_map_and_sen ] +def test_extraction_mode_health_check_resolves_the_mode_from_the_cost_map_and_sends_the_extraction_probe( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + handle: Final = register_scenario(f"health-extraction-{uuid.uuid4().hex[:12]}", _EXTRACTION_PROBE_REPLY) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="cloudflare/@cf/cloudflare/clairic", api_base=handle.api_base(), api_key="synthetic-cloudflare-key" + ) + report: Final = _health_report(gateway, model) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert _probes_sent_to(gateway, handle) == [ + ( + f"/{handle.scenario_id}/ai/run/@cf/cloudflare/clairic", + { + "model": "clairic", + "state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"), + "fields": {"reachable": {"type": "text", "instructions": "Any word from the text"}}, + }, + ) + ] + + @pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation")) def test_decisions_mode_health_check_sends_the_configured_state_and_questions(gateway: Gateway, mode: str) -> None: with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py b/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py new file mode 100644 index 00000000000..7b2a48c9676 --- /dev/null +++ b/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py @@ -0,0 +1,418 @@ +import json +import math +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +Route = Literal["/v1/extraction", "/extraction"] + +_API_KEY: Final = "synthetic-cloudflare-key" +_CLAIRIC: Final = "cloudflare/@cf/cloudflare/clairic" +_CLAIRIC_PATH: Final = "/ai/run/@cf/cloudflare/clairic" +_INPUT_TOKENS: Final = 412 +_EMAIL: Final = "billing@acmetools.example" +_SCENARIO: Final = re.compile(r"scenario-([a-z-]+)") +_JSON_ROWS: Final = TypeAdapter(list[dict[str, JsonValue]]) +_FIELDS: Final[dict[str, JsonValue]] = { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "total": {"type": "number", "unit": "USD"}, +} +_SPEND_QUERY: Final = ( + 'SELECT spend, status, call_type, model_group, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s" +) +_OWNED_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_REQUEST_GUARDRAIL: Final = "extraction-mask-email-in" +_RESPONSE_GUARDRAIL: Final = "extraction-mask-email-out" +_OWNED_MODEL: Final = "invoice-reader" + + +def _state(marker: str, scenario: str = "plain") -> str: + return f"Invoice from Acme Tools, total $174. scenario-{scenario} marker-{marker}" + + +def _vendor_span(state: str) -> list[JsonValue]: + start: Final = state.index("Acme Tools") + return [start, start + len("Acme Tools")] + + +def _email_span(state: str) -> list[JsonValue]: + start: Final = state.index(_EMAIL) + return [start, start + len(_EMAIL)] + + +def _answers(state: str) -> dict[str, JsonValue]: + return { + "vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)}, + "total": {"type": "number", "value": 174, "confidence": 0.95}, + } + + +def _envelope(result: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps({"result": result, "success": True, "errors": [], "messages": []}).encode()) + + +def _clairic_reply(request: Request) -> Reply: + state: Final = JSON_OBJECT.validate_json(request.body)["state"] + assert isinstance(state, str), request.body + return _envelope( + { + "model": "clairic", + "answers": _answers(state), + "truncated": False, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + ) + + +def _expected_response(model: str, state: str) -> dict[str, JsonValue]: + return { + "model": model, + "answers": _answers(state), + "truncated": False, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + + +def _upstream_body(request: Request) -> dict[str, JsonValue]: + assert (request.method, request.target) == ("POST", _CLAIRIC_PATH), (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers + return JSON_OBJECT.validate_json(request.body) + + +def _body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "state": _state(marker), "fields": _FIELDS, **extra} + + +def _extract( + gateway: Gateway, + model: str, + marker: str, + *, + route: Route = "/v1/extraction", + key: str | None = None, + **extra: JsonValue, +) -> httpx.Response: + return gateway.request("POST", route, _body(model, marker, **extra), key=key) + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _expected_spend() -> float: + cost_map: Final = JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) + priced: Final = cost_map.get(_CLAIRIC) + assert priced is not None, f"{_CLAIRIC} has no cost map entry" + return _INPUT_TOKENS * _number(object_value(priced)["input_cost_per_token"]) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _assert_billed(gateway: Gateway, response: httpx.Response, model: str) -> None: + expected: Final = _expected_spend() + assert math.isclose(float(response.headers["x-litellm-response-cost"]), expected, rel_tol=1e-9), response.headers + call_id: Final = response.headers["x-litellm-call-id"] + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"], row["prompt_tokens"], row["completion_tokens"]) == ( + "success", + "aextraction", + model, + _INPUT_TOKENS, + 0, + ), row + assert math.isclose(_number(row["spend"]), expected, rel_tol=1e-9), row + listed: Final = gateway.request("GET", "/spend/logs", params={"request_id": call_id}) + assert listed.status_code == 200, listed.text + (reported,) = _JSON_ROWS.validate_json(listed.content) + assert (reported["request_id"], reported["model_group"]) == (call_id, model), reported + assert math.isclose(_number(reported["spend"]), expected, rel_tol=1e-9), reported + + +@pytest.mark.parametrize("route", ("/v1/extraction", "/extraction")) +def test_extraction_reaches_clairic_on_both_routes_and_is_billed(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = _extract(gateway, model, marker, route=route, metadata={"tags": [f"audit-{marker}"]}) + assert response.status_code == 200, response.text + assert response.json() == _expected_response(model, _state(marker)) + (request,) = wire.drain() + assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS} + _assert_billed(gateway, response, model) + + +_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("unknown_type", {"fields": {"vendor": {"type": "paragraph"}}}), + ("list_without_items", {"fields": {"items": {"type": "list"}}}), + ("remote_image", {"images": ["https://images.example/receipt.png"]}), +) + + +@pytest.mark.parametrize(("label", "change"), _INVALID_BODIES, ids=[label for label, _ in _INVALID_BODIES]) +def test_invalid_extraction_bodies_are_refused_before_any_upstream_call( + gateway: Gateway, label: str, change: dict[str, JsonValue] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/extraction", {**_body(model, marker), **change}) + assert response.status_code == 400, (label, response.text) + assert "Invalid Extraction request" in response.text, (label, response.text) + assert wire.drain() == (), label + + +@pytest.mark.parametrize(("param", "value"), (("temperature", 0.2), ("stream", True)), ids=("temperature", "stream")) +def test_unsupported_top_level_params_are_refused_unless_the_deployment_drops_params( + gateway: Gateway, param: str, value: JsonValue +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + dropping: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY, drop_params=True) + refused: Final = _extract(gateway, strict, marker, **{param: value}) + assert refused.status_code == 400, refused.text + assert param in refused.text, refused.text + assert wire.drain() == () + dropped: Final = _extract(gateway, dropping, marker, **{param: value}) + assert dropped.status_code == 200, dropped.text + assert dropped.headers["content-type"].startswith("application/json"), dropped.headers + assert dropped.json() == _expected_response(dropping, _state(marker)) + (request,) = wire.drain() + assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS} + _assert_billed(gateway, dropped, dropping) + + +def test_key_checks_on_extraction_match_chat(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + anonymous: Final = gateway.client.post("/v1/extraction", json=_body(model, marker)) + assert anonymous.status_code == 401, anonymous.text + restricted: Final = scenario.key(models=[f"other-{uuid.uuid4().hex}"]) + refused: Final = _extract(gateway, model, marker, key=restricted) + chat_refused: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=restricted, + ) + assert 400 <= refused.status_code < 500, refused.text + assert refused.status_code == chat_refused.status_code, (refused.text, chat_refused.text) + assert wire.drain() == () + allowed: Final = scenario.key(models=[model], max_budget=1e-06) + first: Final = _extract(gateway, model, marker, key=allowed) + assert first.status_code == 200, first.text + assert len(wire.drain()) == 1 + blocked, upstream_calls = eventually( + lambda: (_extract(gateway, model, uuid.uuid4().hex, key=allowed), len(wire.drain())), + lambda answer: answer[0].status_code != 200, + seconds=70, + ) + chat_blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=allowed, + ) + assert 400 <= blocked.status_code < 500, blocked.text + assert blocked.status_code == chat_blocked.status_code, (blocked.text, chat_blocked.text) + assert upstream_calls == 0 + assert wire.drain() == () + + +def test_a_small_key_tpm_limit_serves_extraction_and_counts_its_input_tokens(gateway: Gateway) -> None: + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model], tpm_limit=100) + first: Final = _extract(gateway, model, uuid.uuid4().hex, key=key) + assert first.status_code == 200, first.text + assert len(wire.drain()) == 1 + limited, upstream_calls = eventually( + lambda: (_extract(gateway, model, uuid.uuid4().hex, key=key), len(wire.drain())), + lambda answer: answer[0].status_code != 200, + seconds=70, + ) + assert limited.status_code == 429, limited.text + assert upstream_calls == 0 + + +def test_cloudflare_errors_keep_their_status_and_log_an_unbilled_failure(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + refusal: Final = json.dumps( + { + "result": None, + "success": False, + "errors": [{"code": 5006, "message": f"min above max {marker}"}], + "messages": [], + } + ).encode() + with wire_server(lambda _: Reply(status=422, body=refusal)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = _extract(gateway, model, marker) + assert response.status_code == 422, response.text + assert f"min above max {marker}" in response.text, response.text + assert len(wire.drain()) >= 1 + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == ( + "failure", + "aextraction", + model, + 0.0, + ), row + + +def _owned_reply(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + state: Final = body["state"] + assert isinstance(state, str), request.body + scenario: Final = _SCENARIO.search(state) + assert scenario is not None, state + if scenario.group(1) != "mask-out": + return _clairic_reply(request) + return _envelope( + { + "model": "clairic", + "answers": { + "vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)}, + "contact_email": {"type": "text", "value": _EMAIL, "confidence": 0.9, "span": _email_span(state)}, + }, + "truncated": False, + "state_text": state, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + ) + + +def _owned_config(upstream_url: str) -> dict[str, JsonValue]: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + email_filter: Final[list[JsonValue]] = [{"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"}] + return { + **base, + "guardrails": [ + { + "guardrail_name": _REQUEST_GUARDRAIL, + "litellm_params": {"guardrail": "litellm_content_filter", "mode": "pre_call", "patterns": email_filter}, + }, + { + "guardrail_name": _RESPONSE_GUARDRAIL, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "patterns": email_filter, + }, + }, + ], + "model_list": [ + { + "model_name": _OWNED_MODEL, + "litellm_params": {"model": _CLAIRIC, "api_base": upstream_url, "api_key": _API_KEY}, + } + ], + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: OwnedProxy + upstream: Wire + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with wire_server(_owned_reply) as upstream: + directory: Final = tmp_path_factory.mktemp("clairic-extraction-owned") + config_path: Final = directory / "clairic-extraction.yaml" + config_path.write_text(yaml.safe_dump(_owned_config(upstream.url))) + with owned_proxy_process( + session_gateway, + directory, + {}, + config=config_path, + workers=2, + extra_arguments=("--max_tokens", "256", "--temperature", "0.2"), + healthcheck=WorkerHealthcheck.LENIENT, + ) as proxy: + yield Rig(proxy, upstream) + + +def _owned_extract(rig: Rig, state: str, **extra: JsonValue) -> httpx.Response: + return rig.proxy.gateway.request( + "POST", "/v1/extraction", {"model": _OWNED_MODEL, "state": state, "fields": _FIELDS, **extra} + ) + + +def _applied_guardrails(response: httpx.Response) -> frozenset[str]: + return frozenset(name.strip() for name in response.headers.get("x-litellm-applied-guardrails", "").split(",")) + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_cli_generation_defaults_never_reach_an_extraction_call(rig: Rig) -> None: + state: Final = _state(uuid.uuid4().hex) + response: Final = _owned_extract(rig, state) + assert response.status_code == 200, response.text + assert response.json() == _expected_response(_OWNED_MODEL, state) + (request,) = rig.upstream.drain() + assert _upstream_body(request) == {"model": "clairic", "state": state, "fields": _FIELDS} + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_a_request_guardrail_masks_the_state_before_clairic_and_drops_spans_into_it(rig: Rig) -> None: + state: Final = f"{_state(uuid.uuid4().hex, 'mask-in')} Questions go to {_EMAIL}." + response: Final = _owned_extract(rig, state, guardrails=[_REQUEST_GUARDRAIL]) + assert response.status_code == 200, response.text + assert _REQUEST_GUARDRAIL in _applied_guardrails(response), response.headers + (request,) = rig.upstream.drain() + sent: Final = _upstream_body(request) + sent_state: Final = sent["state"] + assert isinstance(sent_state, str) and _EMAIL not in sent_state, sent + assert sent_state.startswith(state.removesuffix(f"{_EMAIL}.")) and "REDACTED" in sent_state, sent + assert sent == {"model": "clairic", "state": sent_state, "fields": _FIELDS} + answers: Final = object_value(response.json()["answers"]) + assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_a_response_guardrail_masks_text_answers_and_state_text(rig: Rig) -> None: + state: Final = f"{_state(uuid.uuid4().hex, 'mask-out')} Questions go to {_EMAIL}." + fields: Final[dict[str, JsonValue]] = {**_FIELDS, "contact_email": {"type": "text"}} + response: Final = rig.proxy.gateway.request( + "POST", + "/v1/extraction", + { + "model": _OWNED_MODEL, + "state": state, + "fields": fields, + "guardrails": [_RESPONSE_GUARDRAIL], + "options": {"return_state_text": True}, + }, + ) + assert response.status_code == 200, response.text + assert _RESPONSE_GUARDRAIL in _applied_guardrails(response), response.headers + (request,) = rig.upstream.drain() + assert _upstream_body(request)["state"] == state + body: Final = response.json() + state_text: Final = body["state_text"] + assert isinstance(state_text, str) and _EMAIL not in state_text and "REDACTED" in state_text, body + answers: Final = object_value(body["answers"]) + contact: Final = object_value(answers["contact_email"]) + assert "span" not in contact, answers + masked: Final = contact["value"] + assert isinstance(masked, str) and _EMAIL not in masked and "REDACTED" in masked, answers + assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers From 93cfc5265eb4df6e8aa463913c684c08728b5561 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 28 Oct 2026 17:44:48 -0700 Subject: [PATCH 5/9] ci: run zizmor with --no-ignores and move the spend logs and cost map guards to pull_request (#48947) * ci(zizmor): run zizmor 1.31.1 without ignores and own the workflow zizmor-action has no input for --no-ignores, so the workflow now installs the zizmor 1.31.1 release binary directly, pinned by SHA-256 like grype in image-scan.yml, and runs it with --no-ignores and --no-config. Inline `zizmor: ignore` comments and config-file ignores no longer hide findings. Behavior otherwise matches the old action defaults (regular persona, online audits with the job token, medium and above, GitHub annotations) Adds a CODEOWNERS entry for .github/workflows/zizmor.yml * ci: run spend-logs-guard and cost-map-guard on pull_request Both guards ran on pull_request_target behind a zizmor ignore comment. They now run on pull_request and check out the base commit by sha, so the guard scripts still come from the base branch while the PR's head is only fetched and diffed. The zizmor ignore comments are gone --- .github/CODEOWNERS | 1 + .github/workflows/cost-map-guard.yml | 5 +++-- .github/workflows/spend-logs-guard.yml | 6 +++--- .github/workflows/zizmor.yml | 17 +++++++++++------ 4 files changed, 18 insertions(+), 11 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f6c06f7d83f..158110cbbab 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,3 +2,4 @@ /litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri /helm/ @yassin-berriai /terraform/ @yassin-berriai +/.github/workflows/zizmor.yml @yuneng-berri diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml index f9172b9551a..05b22017cbb 100644 --- a/.github/workflows/cost-map-guard.yml +++ b/.github/workflows/cost-map-guard.yml @@ -1,7 +1,7 @@ name: Cost map guard -on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed - pull_request_target: +on: + pull_request: branches: - main - "litellm_**" @@ -19,6 +19,7 @@ jobs: steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: + ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - name: Fetch the pull request head and its merge base id: revisions diff --git a/.github/workflows/spend-logs-guard.yml b/.github/workflows/spend-logs-guard.yml index 29e9f236202..cc3bda9b880 100644 --- a/.github/workflows/spend-logs-guard.yml +++ b/.github/workflows/spend-logs-guard.yml @@ -1,7 +1,7 @@ name: Spend logs guard -on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed - pull_request_target: +on: + pull_request: branches: - main @@ -19,7 +19,7 @@ jobs: steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: - ref: ${{ github.event.pull_request.base.ref }} + ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - name: Fetch the pull request head, its merge base and its reviews id: revisions diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index 73f9efb8df9..eba23e795f3 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -27,10 +27,15 @@ jobs: with: persist-credentials: false + - name: Install zizmor + run: | + curl -fsSL --retry 3 -o "$RUNNER_TEMP/zizmor.tar.gz" \ + https://github.com/zizmorcore/zizmor/releases/download/v1.31.1/zizmor-x86_64-unknown-linux-gnu.tar.gz + echo "d185f726bdaf92df78273217dfdd24ebcd6b0ef48cf3ce776ff5762b6e138c99 $RUNNER_TEMP/zizmor.tar.gz" | sha256sum -c - + tar xzf "$RUNNER_TEMP/zizmor.tar.gz" -C "$RUNNER_TEMP" zizmor + - name: Run zizmor - uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6 - with: - version: "1.24.1" - min-severity: medium - advanced-security: false - annotations: true + env: + GH_TOKEN: ${{ github.token }} + run: | + "$RUNNER_TEMP/zizmor" --no-ignores --no-config --min-severity=medium --format=github . From e1580b7b837b680d0a0770cfa5f6c88938bc6050 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 17:48:09 -0700 Subject: [PATCH 6/9] refactor(images): remove dall-e-2 and dall-e-3 code paths (#48703) OpenAI shut down dall-e-2 and dall-e-3 on 2026-05-12. Drop their image generation and edit configs, the Azure dall-e-2 create-and-poll flow, and the silent dall-e-2 fallback when no model is passed. image_generation without a model or provider now fails locally with "model is required". image_edit and image_variation leave model out of the provider request instead of sending dall-e-2, so providers that ignore model (Topaz) keep working and OpenAI returns its own 400. --- litellm/__init__.py | 1 - litellm/constants.py | 1 - litellm/cost_calculator.py | 11 +- litellm/images/main.py | 32 ++- .../get_llm_provider_logic.py | 1 - litellm/llms/azure/azure.py | 196 +----------------- .../llms/azure/image_generation/__init__.py | 10 +- .../dall_e_2_transformation.py | 7 - .../dall_e_3_transformation.py | 7 - .../azure_ai/image_generation/__init__.py | 10 +- .../dall_e_2_transformation.py | 7 - .../dall_e_3_transformation.py | 7 - litellm/llms/custom_httpx/aiohttp_handler.py | 5 +- litellm/llms/openai/image_edit/__init__.py | 26 --- .../image_edit/dalle2_transformation.py | 98 --------- .../llms/openai/image_edit/transformation.py | 2 +- .../llms/openai/image_generation/__init__.py | 17 -- .../dall_e_2_transformation.py | 81 -------- .../dall_e_3_transformation.py | 81 -------- .../llms/openai/image_variations/handler.py | 2 +- litellm/utils.py | 11 +- tests/unit/images/test_image_edit.py | 18 ++ tests/unit/images/test_main.py | 23 +- .../test_azure_image_generation_init.py | 18 -- ...t_openai_image_generation_extra_headers.py | 2 +- .../llms/openai/test_openai_common_utils.py | 38 ++-- 26 files changed, 89 insertions(+), 623 deletions(-) delete mode 100644 litellm/llms/azure/image_generation/dall_e_2_transformation.py delete mode 100644 litellm/llms/azure/image_generation/dall_e_3_transformation.py delete mode 100644 litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py delete mode 100644 litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py delete mode 100644 litellm/llms/openai/image_edit/dalle2_transformation.py delete mode 100644 litellm/llms/openai/image_generation/dall_e_2_transformation.py delete mode 100644 litellm/llms/openai/image_generation/dall_e_3_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 8dfe5adca1c..c649a48211a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1336,7 +1336,6 @@ all_embedding_models = ( ) ####### IMAGE GENERATION MODELS ################### -openai_image_generation_models = ["dall-e-2", "dall-e-3"] ####### VIDEO GENERATION MODELS ################### openai_video_generation_models = ["sora-2"] diff --git a/litellm/constants.py b/litellm/constants.py index 1d5c6404bfd..25770d960bb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -764,7 +764,6 @@ ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { # Used for web search interception across providers LITELLM_WEB_SEARCH_TOOL_NAME: Final = "litellm_web_search" -DEFAULT_IMAGE_ENDPOINT_MODEL: Final = "dall-e-2" DEFAULT_VIDEO_ENDPOINT_MODEL: Final = "sora-2" DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = get_env_int("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 86b22069a39..087d5b5b718 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1592,14 +1592,6 @@ def completion_cost( for part in ws_tier_parts ) - if ( - (call_type == "aimage_generation" or call_type == "image_generation") - and model is not None - and isinstance(model, str) - and len(model) == 0 - and custom_llm_provider == "azure" - ): - model = "dall-e-2" # for dall-e-2, azure expects an empty model name # Handle Inputs to completion_cost prompt_tokens = 0 prompt_characters: int | None = None @@ -1674,7 +1666,8 @@ def completion_cost( raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates) ) - for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names): + for idx, (candidate_model, candidate_pricing_provider) in enumerate(potential_model_names): + model = candidate_model try: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("selected model name for cost calculation: %s", model) diff --git a/litellm/images/main.py b/litellm/images/main.py index d0618e5e804..6e77cbb7046 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -16,7 +16,6 @@ import litellm # client is imported from litellm as it's a decorator from litellm import client -from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.exceptions import LiteLLMUnknownProvider from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -299,17 +298,14 @@ def image_generation( if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/")) else None ) - dynamic_api_key: str | None = None - if model is not None or custom_llm_provider is not None: - model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - litellm_params=chatgpt_params, - ) - else: - model = "dall-e-2" - custom_llm_provider = "openai" # default to dall-e-2 on openai + if model is None and custom_llm_provider is None: + raise litellm.BadRequestError(message="model is required", model="", llm_provider="") + model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + litellm_params=chatgpt_params, + ) model_response.hidden_params["model"] = model openai_params: Final = [ "user", @@ -712,7 +708,7 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: @client def image_variation( image: FileTypes, - model: str = "dall-e-2", # set to dall-e-2 by default - like OpenAI. + model: str | None = None, n: int = 1, response_format: Literal["url", "b64_json"] = "url", size: str | None = None, @@ -728,7 +724,7 @@ def image_variation( litellm_params: Final = get_litellm_params(**kwargs) # get the custom llm provider model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, + model=model or "", custom_llm_provider=litellm_params.get("custom_llm_provider", None), api_base=litellm_params.get("api_base", None), api_key=litellm_params.get("api_key", None), @@ -747,7 +743,7 @@ def image_variation( response: ImageResponse | None = None provider_config: Final = ProviderConfigManager.get_provider_model_info( - model=model or "", # openai defaults to dall-e-2 + model=model, provider=llm_provider, ) @@ -772,7 +768,7 @@ def image_variation( model_response=model_response, api_key=api_key, api_base=api_base, - model=model, + model=model or None, image=image, timeout=litellm_params.get("timeout", None), custom_llm_provider=custom_llm_provider, @@ -790,7 +786,7 @@ def image_variation( model_response=model_response, api_key=api_key, api_base=api_base, - model=model, + model=model or None, image=image, timeout=litellm_params.get("timeout", None) or DEFAULT_REQUEST_TIMEOUT, custom_llm_provider=custom_llm_provider, @@ -901,7 +897,7 @@ def image_edit( # get llm provider logic litellm_params: Final = GenericLiteLLMParams(**kwargs) model, custom_llm_provider, _, _ = get_llm_provider( - model=model or DEFAULT_IMAGE_ENDPOINT_MODEL, + model=model or "", custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 19c1d5c414c..d9933087334 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -437,7 +437,6 @@ def get_llm_provider( model in litellm.open_ai_chat_completion_models or "ft:gpt-3.5-turbo" in model or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o - or model in litellm.openai_image_generation_models or model.startswith("gpt-image") or model in litellm.openai_video_generation_models ): diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 02f1ba3472a..5a181faad11 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -13,13 +13,12 @@ from openai import ( BadRequestError, OpenAI, ) -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel import litellm -from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES +from litellm.constants import DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing -from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -34,7 +33,6 @@ from litellm.types.utils import ( from litellm.utils import ( CustomStreamWrapper, convert_to_model_response_object, - modify_url, ) from ...types.llms.openai import HttpxBinaryResponseContent @@ -59,8 +57,6 @@ from .image_generation import ( ) from .image_generation.http_utils import azure_deployment_image_generation_json_body -_DECODED_JSON: Final = TypeAdapter(object) - class AzureOpenAIAssistantsAPIConfig: """ @@ -860,11 +856,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, deployment_name: str | None = None, ) -> httpx.Response: - """ - Implemented for azure dall-e-2 image gen calls - - Alternative to needing a custom transport implementation - """ if client is None: _params: Final = {} if timeout is not None: @@ -881,95 +872,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): else: async_handler = client - if ( - "images/generations" in api_base - and api_version - in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict - "2023-06-01-preview", - "2023-07-01-preview", - "2023-08-01-preview", - "2023-09-01-preview", - "2023-10-01-preview", - ] - ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") - - data.pop( - "model", None - ) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview - response = await async_handler.post( - url=api_base, - data=json.dumps(data), - headers=headers, - ) - if "operation-location" in response.headers: - operation_location_url: Final = response.headers["operation-location"] - else: - raise AzureOpenAIError(status_code=500, message=response.text) - # Reject polling URLs that don't share an origin with ``api_base``. - # Without this an upstream-controlled or attacker-controlled - # value would receive the operator's Azure API key in the - # request headers below. VERIA-51. - try: - assert_same_origin(operation_location_url, api_base) - except SSRFError as ssrf_err: - raise AzureOpenAIError( - status_code=502, - message=f"Rejected polling URL: {ssrf_err}", - ) - response = await async_handler.get( - url=operation_location_url, - headers=headers, - ) - - await response.aread() - - timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT - start_time: Final = time.time() - if "status" not in response.json(): - # Don't reflect the raw response body — when the polling - # URL points at an internal JSON API (cloud metadata - # service etc.) reflecting it here turns Blind SSRF into - # Full-Read SSRF. VERIA-51. - raise AzureOpenAIError( - status_code=502, - message="Polling response missing 'status' field", - ) - while response.json()["status"] not in ["succeeded", "failed"]: - if time.time() - start_time > timeout_secs: - raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") - - await asyncio.sleep(int(response.headers.get("retry-after") or 10)) - response = await async_handler.get( - url=operation_location_url, - headers=headers, - ) - await response.aread() - - if response.json()["status"] == "failed": - error_data: Final = response.json() - # Preserve Azure error details (e.g. content_policy_violation, - # inner_error, content_filter_results) as structured body so - # exception_type() can route them correctly. - _error_body: Final = error_data.get("error", error_data) - _error_msg: Final = ( - _error_body.get("message", "Image generation failed") - if isinstance(_error_body, dict) - else json.dumps(error_data) - ) - raise AzureOpenAIError( - status_code=400, - message=_error_msg, - body=error_data, - ) - - result: Final = _DECODED_JSON.validate_python(response.json()["result"]) - return httpx.Response( - status_code=200, - headers=response.headers, - content=json.dumps(result).encode("utf-8"), - request=httpx.Request(method="POST", url="https://api.openai.com/v1"), - ) request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return await async_handler.post( url=api_base, @@ -988,11 +890,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, deployment_name: str | None = None, ) -> httpx.Response: - """ - Implemented for azure dall-e-2 image gen calls - - Alternative to needing a custom transport implementation - """ if client is None: _params: Final = {} if timeout is not None: @@ -1006,87 +903,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): else: sync_handler = client - if ( - "images/generations" in api_base - and api_version - in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict - "2023-06-01-preview", - "2023-07-01-preview", - "2023-08-01-preview", - "2023-09-01-preview", - "2023-10-01-preview", - ] - ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") - - data.pop( - "model", None - ) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview - response = sync_handler.post( - url=api_base, - data=json.dumps(data), - headers=headers, - ) - if "operation-location" in response.headers: - operation_location_url: Final = response.headers["operation-location"] - else: - raise AzureOpenAIError(status_code=500, message=response.text) - try: - assert_same_origin(operation_location_url, api_base) - except SSRFError as ssrf_err: - raise AzureOpenAIError( - status_code=502, - message=f"Rejected polling URL: {ssrf_err}", - ) - response = sync_handler.get( - url=operation_location_url, - headers=headers, - ) - - response.read() - - timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT - start_time: Final = time.time() - if "status" not in response.json(): - raise AzureOpenAIError( - status_code=502, - message="Polling response missing 'status' field", - ) - while response.json()["status"] not in ["succeeded", "failed"]: - if time.time() - start_time > timeout_secs: - raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") - - time.sleep(int(response.headers.get("retry-after") or 10)) - response = sync_handler.get( - url=operation_location_url, - headers=headers, - ) - response.read() - - if response.json()["status"] == "failed": - error_data: Final = response.json() - # Preserve Azure error details (e.g. content_policy_violation, - # inner_error, content_filter_results) as structured body so - # exception_type() can route them correctly. - _error_body: Final = error_data.get("error", error_data) - _error_msg: Final = ( - _error_body.get("message", "Image generation failed") - if isinstance(_error_body, dict) - else json.dumps(error_data) - ) - raise AzureOpenAIError( - status_code=400, - message=_error_msg, - body=error_data, - ) - - result: Final = _DECODED_JSON.validate_python(response.json()["result"]) - return httpx.Response( - status_code=200, - headers=response.headers, - content=json.dumps(result).encode("utf-8"), - request=httpx.Request(method="POST", url="https://api.openai.com/v1"), - ) request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return sync_handler.post( url=api_base, @@ -1193,10 +1009,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): deployment_name=model, ) - provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) + provider_config: Final = get_azure_image_generation_config(data.get("model", "")) if provider_config is not None: return provider_config.transform_image_generation_response( - model=data.get("model", "dall-e-2"), + model=data.get("model", ""), raw_response=httpx_response, model_response=model_response or ImageResponse(), logging_obj=logging_obj, @@ -1329,10 +1145,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers=request_headers, deployment_name=model, ) - provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) + provider_config: Final = get_azure_image_generation_config(data.get("model", "")) if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): return provider_config.transform_image_generation_response( - model=data.get("model", "dall-e-2"), + model=data.get("model", ""), raw_response=httpx_response, model_response=model_response or ImageResponse(), logging_obj=logging_obj, diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index a2a905f2287..7e9b5f5d99c 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -4,14 +4,10 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig -from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig from .gpt_transformation import AzureGPTImageGenerationConfig from .http_utils import azure_deployment_image_generation_json_body __all__ = [ - "AzureDallE2ImageGenerationConfig", - "AzureDallE3ImageGenerationConfig", "AzureGPTImageGenerationConfig", "azure_deployment_image_generation_json_body", ] @@ -21,11 +17,7 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: model = model.lower() model = model.replace("-", "") model = model.replace("_", "") - if model == "" or "dalle2" in model: # empty model is dall-e-2 - return AzureDallE2ImageGenerationConfig() - elif "dalle3" in model: - return AzureDallE3ImageGenerationConfig() - elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): return AzureFoundryMAIImageGenerationConfig() else: verbose_logger.debug( diff --git a/litellm/llms/azure/image_generation/dall_e_2_transformation.py b/litellm/llms/azure/image_generation/dall_e_2_transformation.py deleted file mode 100644 index 788753e4ea9..00000000000 --- a/litellm/llms/azure/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig - - -class AzureDallE2ImageGenerationConfig(DallE2ImageGenerationConfig): - """ - Azure dall-e-2 image generation config - """ diff --git a/litellm/llms/azure/image_generation/dall_e_3_transformation.py b/litellm/llms/azure/image_generation/dall_e_3_transformation.py deleted file mode 100644 index e6974f9317e..00000000000 --- a/litellm/llms/azure/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig - - -class AzureDallE3ImageGenerationConfig(DallE3ImageGenerationConfig): - """ - Azure dall-e-3 image generation config - """ diff --git a/litellm/llms/azure_ai/image_generation/__init__.py b/litellm/llms/azure_ai/image_generation/__init__.py index 742922ad295..bd3c11d36f9 100644 --- a/litellm/llms/azure_ai/image_generation/__init__.py +++ b/litellm/llms/azure_ai/image_generation/__init__.py @@ -3,15 +3,11 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from .dall_e_2_transformation import AzureFoundryDallE2ImageGenerationConfig -from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig from .flux_transformation import AzureFoundryFluxImageGenerationConfig from .gpt_transformation import AzureFoundryGPTImageGenerationConfig from .mai_transformation import AzureFoundryMAIImageGenerationConfig __all__ = [ - "AzureFoundryDallE2ImageGenerationConfig", - "AzureFoundryDallE3ImageGenerationConfig", "AzureFoundryFluxImageGenerationConfig", "AzureFoundryGPTImageGenerationConfig", "AzureFoundryMAIImageGenerationConfig", @@ -22,11 +18,7 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi model = model.lower() model = model.replace("-", "") model = model.replace("_", "") - if model == "" or "dalle2" in model: # empty model is dall-e-2 - return AzureFoundryDallE2ImageGenerationConfig() - elif "dalle3" in model: - return AzureFoundryDallE3ImageGenerationConfig() - elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): return AzureFoundryMAIImageGenerationConfig() elif "flux" in model: return AzureFoundryFluxImageGenerationConfig() diff --git a/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py b/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py deleted file mode 100644 index 2c8a5426628..00000000000 --- a/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig - - -class AzureFoundryDallE2ImageGenerationConfig(DallE2ImageGenerationConfig): - """ - Azure dall-e-2 image generation config - """ diff --git a/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py b/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py deleted file mode 100644 index 3ceb9b405da..00000000000 --- a/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig - - -class AzureFoundryDallE3ImageGenerationConfig(DallE3ImageGenerationConfig): - """ - Azure dall-e-3 image generation config - """ diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 4060d84f7f4..e6544968d24 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -561,11 +561,8 @@ class BaseLLMAIOHTTPHandler: organization: str | None = None, headers: dict | None = None, ) -> ImageResponse: - if model is None: - raise ValueError("model is required for non-openai image variations") - provider_config: Final = ProviderConfigManager.get_provider_image_variation_config( - model=model, # openai defaults to dall-e-2 + model=model or "", provider=LlmProviders(custom_llm_provider), ) diff --git a/litellm/llms/openai/image_edit/__init__.py b/litellm/llms/openai/image_edit/__init__.py index 63a41b65e8a..261010071c9 100644 --- a/litellm/llms/openai/image_edit/__init__.py +++ b/litellm/llms/openai/image_edit/__init__.py @@ -1,31 +1,5 @@ -from typing import Final - -from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig - -from .dalle2_transformation import DallE2ImageEditConfig from .transformation import OpenAIImageEditConfig __all__ = [ - "DallE2ImageEditConfig", "OpenAIImageEditConfig", - "get_openai_image_edit_config", ] - - -def get_openai_image_edit_config(model: str) -> BaseImageEditConfig: - """ - Get the appropriate OpenAI image edit config based on the model. - - Args: - model: The model name (e.g., "dall-e-2", "gpt-image-1") - - Returns: - The appropriate config instance for the model - """ - model_normalized: Final = model.lower().replace("-", "").replace("_", "") - - if model_normalized == "dalle2": - return DallE2ImageEditConfig() - else: - # Default to standard OpenAI config for gpt-image-1 and other models - return OpenAIImageEditConfig() diff --git a/litellm/llms/openai/image_edit/dalle2_transformation.py b/litellm/llms/openai/image_edit/dalle2_transformation.py deleted file mode 100644 index f349959b14f..00000000000 --- a/litellm/llms/openai/image_edit/dalle2_transformation.py +++ /dev/null @@ -1,98 +0,0 @@ -from io import BufferedReader -from typing import TYPE_CHECKING, Any, Final, cast - -from httpx._types import RequestFiles - -import litellm -from litellm.images.utils import ImageEditRequestUtils -from litellm.types.images.main import ImageEditRequestParams -from litellm.types.llms.openai import FileTypes -from litellm.types.router import GenericLiteLLMParams - -from .transformation import OpenAIImageEditConfig - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - - LiteLLMLoggingObj = _LiteLLMLoggingObj -else: - LiteLLMLoggingObj = Any - - -class DallE2ImageEditConfig(OpenAIImageEditConfig): - """ - DALL-E-2 specific configuration for image edit API. - - DALL-E-2 only supports editing a single image (not an array). - Uses "image" field name instead of "image[]". - """ - - def transform_image_edit_request( - self, - model: str, - prompt: str | None, - image: FileTypes | None, - image_edit_optional_request_params: dict, - litellm_params: GenericLiteLLMParams, - headers: dict, - ) -> tuple[dict, RequestFiles]: - """ - Transform image edit request for DALL-E-2. - - DALL-E-2 only accepts a single image with field name "image" (not "image[]"). - """ - request_params: Final = { - "model": model, - **image_edit_optional_request_params, - } - if image is not None: - request_params["image"] = image - if prompt is not None: - request_params["prompt"] = prompt - - request: Final = ImageEditRequestParams(**request_params) - request_dict: Final = cast(dict, request) - - ######################################################### - # Separate images and masks as `files` and send other parameters as `data` - ######################################################### - _image_list: Final = request_dict.get("image") - _mask = request_dict.get("mask") - data_without_files: Final = {k: v for k, v in request_dict.items() if k not in ["image", "mask"]} - files_list: Final[list[tuple[str, Any]]] = [] - - # Handle image parameter - DALL-E-2 only supports single image - if _image_list is not None: - image_list: Final = [_image_list] if not isinstance(_image_list, list) else _image_list - - # Validate only one image is provided - if len(image_list) > 1: - raise litellm.BadRequestError( - message="DALL-E-2 only supports editing a single image. Please provide one image.", - model=model, - llm_provider="openai", - ) - - # Use "image" field name (singular) for DALL-E-2 - for _image in image_list: - if _image is not None: - self._add_image_to_files( - files_list=files_list, - image=_image, - field_name="image", - ) - - # Handle mask parameter if provided - if _mask is not None: - # Handle case where mask can be a list (extract first mask) - if isinstance(_mask, list): - _mask = _mask[0] if _mask else None - - if _mask is not None: - mask_content_type: Final[str] = ImageEditRequestUtils.get_image_content_type(_mask) - if isinstance(_mask, BufferedReader): - files_list.append(("mask", (_mask.name, _mask, mask_content_type))) - else: - files_list.append(("mask", ("mask.png", _mask, mask_content_type))) - - return data_without_files, files_list diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index 9def833052b..0bba58096ef 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -96,7 +96,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): """ # Build request params, only including non-None values request_params: Final = { - "model": model, + **({"model": model} if model else {}), **image_edit_optional_request_params, } if image is not None: diff --git a/litellm/llms/openai/image_generation/__init__.py b/litellm/llms/openai/image_generation/__init__.py index e20c80f20bb..4cffbe81a0e 100644 --- a/litellm/llms/openai/image_generation/__init__.py +++ b/litellm/llms/openai/image_generation/__init__.py @@ -1,9 +1,3 @@ -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) - -from .dall_e_2_transformation import DallE2ImageGenerationConfig -from .dall_e_3_transformation import DallE3ImageGenerationConfig from .gpt_transformation import GPTImageGenerationConfig from .guardrail_translation import ( OpenAIImageGenerationHandler, @@ -11,18 +5,7 @@ from .guardrail_translation import ( ) __all__ = [ - "DallE2ImageGenerationConfig", - "DallE3ImageGenerationConfig", "GPTImageGenerationConfig", "OpenAIImageGenerationHandler", "guardrail_translation_mappings", ] - - -def get_openai_image_generation_config(model: str) -> BaseImageGenerationConfig: - if model.startswith("dall-e-2") or model == "": # empty model is dall-e-2 - return DallE2ImageGenerationConfig() - elif model.startswith("dall-e-3"): - return DallE3ImageGenerationConfig() - else: - return GPTImageGenerationConfig() diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py deleted file mode 100644 index ffc6f1d5fe9..00000000000 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import TYPE_CHECKING, Final - -import httpx - -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageResponse -from litellm.utils import convert_to_model_response_object - -if TYPE_CHECKING: - from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj - - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - - -class DallE2ImageGenerationConfig(BaseImageGenerationConfig): - """ - OpenAI dall-e-2 image generation config - """ - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - return ["n", "response_format", "quality", "size", "user"] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_params: Final = self.get_supported_openai_params(model) - for k in non_default_params: - if k not in optional_params: - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) - - return optional_params - - def transform_image_generation_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ImageResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict, - optional_params: dict, - litellm_params: dict, - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ImageResponse: - response: Final = raw_response.json() - - stringified_response: Final = response - ## LOGGING - logging_obj.post_call( - input=request_data.get("prompt", ""), - api_key=api_key, - additional_args={"complete_input_dict": request_data}, - original_response=stringified_response, - ) - image_response: Final[ImageResponse] = convert_to_model_response_object( - response_object=stringified_response, - model_response_object=model_response, - response_type="image_generation", - ) - - # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "standard") # always standard for dall-e-2 - image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-2 - - return image_response diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py deleted file mode 100644 index 90b7eaedf2f..00000000000 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import TYPE_CHECKING, Final - -import httpx - -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageResponse -from litellm.utils import convert_to_model_response_object - -if TYPE_CHECKING: - from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj - - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - - -class DallE3ImageGenerationConfig(BaseImageGenerationConfig): - """ - OpenAI dall-e-3 image generation config - """ - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - return ["n", "response_format", "quality", "size", "user", "style"] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_params: Final = self.get_supported_openai_params(model) - for k in non_default_params: - if k not in optional_params: - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) - - return optional_params - - def transform_image_generation_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ImageResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict, - optional_params: dict, - litellm_params: dict, - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ImageResponse: - response: Final = raw_response.json() - - stringified_response: Final = response - ## LOGGING - logging_obj.post_call( - input=request_data.get("prompt", ""), - api_key=api_key, - additional_args={"complete_input_dict": request_data}, - original_response=stringified_response, - ) - image_response: Final[ImageResponse] = convert_to_model_response_object( - response_object=stringified_response, - model_response_object=model_response, - response_type="image_generation", - ) - - # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "hd") # always hd for dall-e-3 - image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-3 - - return image_response diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 89dcdc40e02..01f01d3de1d 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -151,7 +151,7 @@ class OpenAIImageVariationsHandler: ) -> ImageResponse: try: provider_config: Final = ProviderConfigManager.get_provider_image_variation_config( - model=model or "", # openai defaults to dall-e-2 + model=model or "", provider=LlmProviders.OPENAI, ) diff --git a/litellm/utils.py b/litellm/utils.py index 1ef39059e1b..be1e11dd39f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7664,7 +7664,6 @@ def validate_environment( model in litellm.open_ai_chat_completion_models or model in litellm.open_ai_text_completion_models or model in litellm.open_ai_embedding_models - or model in litellm.openai_image_generation_models or model.startswith("gpt-image") ): if "OPENAI_API_KEY" in os.environ: @@ -10311,11 +10310,9 @@ class ProviderConfigManager: provider: LlmProviders, ) -> BaseImageGenerationConfig | None: if LlmProviders.OPENAI == provider: - from litellm.llms.openai.image_generation import ( - get_openai_image_generation_config, - ) + from litellm.llms.openai.image_generation import GPTImageGenerationConfig - return get_openai_image_generation_config(model) + return GPTImageGenerationConfig() elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_generation import ( get_azure_image_generation_config, @@ -10546,9 +10543,9 @@ class ProviderConfigManager: provider: LlmProviders, ) -> BaseImageEditConfig | None: if LlmProviders.OPENAI == provider: - from litellm.llms.openai.image_edit import get_openai_image_edit_config + from litellm.llms.openai.image_edit import OpenAIImageEditConfig - return get_openai_image_edit_config(model=model) + return OpenAIImageEditConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config diff --git a/tests/unit/images/test_image_edit.py b/tests/unit/images/test_image_edit.py index fb7a2869b2a..4f858050e3e 100644 --- a/tests/unit/images/test_image_edit.py +++ b/tests/unit/images/test_image_edit.py @@ -105,6 +105,24 @@ def test_openai_image_edit_sync_accepts_bytesio_images(respx_mock: respx.MockRou assert _multipart_image_parts(route.calls[0].request) == (_FIRST_IMAGE,) +def test_openai_image_edit_without_a_model_leaves_model_out_of_the_request( + respx_mock: respx.MockRouter, httpx_transport: None +) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( + return_value=httpx.Response(200, json=_EDIT_RESPONSE) + ) + + litellm.image_edit( + prompt="combine the reference images", + image=io.BytesIO(_FIRST_IMAGE), + custom_llm_provider="openai", + api_key="fake-key", + ) + + assert route.call_count == 1 + assert b'name="model"' not in route.calls[0].request.read() + + @pytest.mark.asyncio async def test_openai_image_edit_accepts_bytesio_images(respx_mock: respx.MockRouter, httpx_transport: None) -> None: route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 29424e50079..61b2f28b727 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -300,7 +300,7 @@ async def test_aimage_generation_ignores_a_non_string_base_url_like_the_sync_pat ) response: Final = await litellm.aimage_generation( - model="dall-e-3", + model="gpt-image-2.5-sunburst", prompt="a red circle", api_base=api_base, base_url=httpx.URL(api_base), @@ -688,3 +688,24 @@ async def test_chatgpt_images_forward_client_login_and_workspace( assert "cookie" not in request.headers assert "x-litellm-api-key" not in request.headers assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"} + + +def test_image_generation_without_a_model_or_provider_fails_before_calling_any_provider( + respx_mock: respx.MockRouter, +) -> None: + with pytest.raises(litellm.BadRequestError, match="model is required"): + litellm.image_generation(prompt="a red circle", api_key="fake_openai_api_key") + + assert not respx_mock.calls + + +def test_topaz_image_variation_works_without_a_model(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://api.topazlabs.com/image/v1/enhance").mock( + return_value=httpx.Response(200, content=b"enhanced-image", headers={"content-type": "image/png"}) + ) + + response: Final = litellm.image_variation(image=b"source-image", custom_llm_provider="topaz", api_key="fake-key") + + assert route.call_count == 1 + assert response.data is not None + assert response.data[0].b64_json == base64.b64encode(b"enhanced-image").decode() diff --git a/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py index eabd5c8427d..83e01fce642 100644 --- a/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py @@ -20,27 +20,9 @@ from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.azure.image_generation import ( - AzureDallE3ImageGenerationConfig, - get_azure_image_generation_config, -) from litellm.utils import get_optional_params_image_gen -@pytest.mark.parametrize( - "received_model, expected_config", - [ - ("dall-e-3", AzureDallE3ImageGenerationConfig), - ("dalle-3", AzureDallE3ImageGenerationConfig), - ("openai_dall_e_3", AzureDallE3ImageGenerationConfig), - ], -) -def test_azure_image_generation_config(received_model, expected_config): - assert isinstance( - get_azure_image_generation_config(received_model), expected_config - ) - - def test_azure_deployment_image_generation_json_body(): """Deployment-scoped Azure image URL must not send ``model`` in JSON.""" api = ( diff --git a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index 11df07f5fea..e103412e0c3 100644 --- a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -265,7 +265,7 @@ class TestImageGenerationEntryPointHeaders: test_headers = {"cf-aig-authorization": "Bearer my-secret"} await litellm.aimage_generation( - model="dall-e-3", + model="gpt-image-2.5-sunburst", prompt="A white cat", extra_headers=test_headers, client=mock_openai_client, diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 831df30c98a..8dc92e23fdf 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -46,12 +46,12 @@ API_FUNCTION_PARAMS = [ ( "image_generation", False, - {"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"}, + {"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"}, ), ( "image_generation", True, - {"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"}, + {"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"}, ), ( "speech", @@ -84,11 +84,7 @@ async def test_openai_client_reuse(function_name, is_async, args): """ # Determine which client class to mock based on whether the test is async - client_path = ( - "litellm.llms.openai.openai.AsyncOpenAI" - if is_async - else "litellm.llms.openai.openai.OpenAI" - ) + client_path = "litellm.llms.openai.openai.AsyncOpenAI" if is_async else "litellm.llms.openai.openai.OpenAI" # Create the appropriate patches with ( @@ -98,9 +94,7 @@ async def test_openai_client_reuse(function_name, is_async, args): ): # Setup the mock to return None first time (cache miss) then a client for subsequent calls mock_client = MagicMock() - mock_get_cache.side_effect = [None] + [ - mock_client - ] * 9 # First call returns None, rest return the mock client + mock_get_cache.side_effect = [None] + [mock_client] * 9 # First call returns None, rest return the mock client # Make 10 API calls for _ in range(10): @@ -118,9 +112,9 @@ async def test_openai_client_reuse(function_name, is_async, args): pass # Verify client was created only once - assert ( - mock_client_class.call_count == 1 - ), f"{'Async' if is_async else ''}OpenAI client should be created only once" + assert mock_client_class.call_count == 1, ( + f"{'Async' if is_async else ''}OpenAI client should be created only once" + ) # Verify the client was cached assert mock_set_cache.call_count == 1, "Client should be cached once" @@ -144,12 +138,8 @@ def test_precomputed_init_params_match_inspect_signature(): _OPENAI_INIT_PARAMS, ) - expected_openai = tuple( - p for p in inspect.signature(OpenAI.__init__).parameters if p != "self" - ) - expected_azure = tuple( - p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self" - ) + expected_openai = tuple(p for p in inspect.signature(OpenAI.__init__).parameters if p != "self") + expected_azure = tuple(p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self") assert _OPENAI_INIT_PARAMS == expected_openai assert _AZURE_OPENAI_INIT_PARAMS == expected_azure @@ -500,7 +490,10 @@ async def test_sdk_api_factories_keep_httpx_transport_and_request_timeouts(api: moved: Final = mock.get(host="sdk-default.example", path="/v1/moved-models").mock( return_value=httpx.Response( 200, - json={"object": "list", "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}]}, + json={ + "object": "list", + "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}], + }, ) ) sdk_client: Final = _sdk_api_client(api, is_async, timeout) @@ -542,7 +535,10 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session( assert azure_client is not None assert azure_client._client is session response: Final = await aimage_variation( - image=BytesIO(b"image-bytes"), api_key="transport-only", api_base="https://sdk-default.example/v1" + image=BytesIO(b"image-bytes"), + custom_llm_provider="openai", + api_key="transport-only", + api_base="https://sdk-default.example/v1", ) assert response.data[0].url == "https://example.com/image.png" assert not session.is_closed From 77cad6e9032369fab248870d8ca9cfa94367484a Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:05:03 -0700 Subject: [PATCH 7/9] test(proxy): audit #46997 (#48918) * test(proxy): audit #46997 Add request-path cells for the global proxy spend counter: a key with no user spends toward litellm max_budget, and two proxies sharing Redis add every request to one counter that matches the proxy budget row * test(proxy): wait for the recorded budget spend and shared counter before checking them --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../spend/test_global_proxy_spend_counter.py | 159 +++++++++++++++++- 1 file changed, 158 insertions(+), 1 deletion(-) diff --git a/tests/integration/spend/test_global_proxy_spend_counter.py b/tests/integration/spend/test_global_proxy_spend_counter.py index 0185e48ffbd..e818487cf53 100644 --- a/tests/integration/spend/test_global_proxy_spend_counter.py +++ b/tests/integration/spend/test_global_proxy_spend_counter.py @@ -1,16 +1,113 @@ +import json import os import uuid +from itertools import product +from pathlib import Path from typing import Final +import httpx import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.redis_process import owned_redis +from pydantic import JsonValue, TypeAdapter +from redis import Redis import litellm from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY +from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import update_cache +BUDGET_MODEL: Final = "proxy-budget-model" +CONFIG: Final = TypeAdapter(dict[str, JsonValue]) +SECTION: Final = TypeAdapter(dict[str, JsonValue]) +OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 120 + + +def _proxy_budget_config(gateway: Gateway, directory: Path, max_budget: float) -> Path: + loaded: Final = CONFIG.validate_python( + yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text()) + ) + model: Final[JsonValue] = { + "model_name": BUDGET_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": f"{gateway.upstream_url}/v1", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + } + path: Final = directory / "proxy_budget_config.yaml" + path.write_text( + yaml.safe_dump( + { + **loaded, + "model_list": [model], + "general_settings": { + **SECTION.validate_python(loaded["general_settings"]), + "user_api_key_cache_ttl": 600, + }, + "litellm_settings": { + **SECTION.validate_python(loaded["litellm_settings"]), + "max_budget": max_budget, + "budget_duration": "30d", + }, + } + ) + ) + return path + + +def _chat(proxy: Gateway, key: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": BUDGET_MODEL, "max_tokens": 20, "messages": [{"role": "user", "content": marker}]}, + key=key, + ) + + +def _served_id(proxy: Gateway, key: str) -> str: + response: Final = _chat(proxy, key, f"proxy budget {uuid.uuid4().hex}") + assert response.status_code == 200, response.text + identity: Final = response.json()["id"] + assert isinstance(identity, str), response.text + return identity + + +def _logged_spend(proxy: Gateway, request_ids: tuple[str, ...]) -> tuple[float, ...]: + rows: Final = eventually( + lambda: proxy.rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', + ("{" + ",".join(request_ids) + "}",), + ), + lambda found: len(found) == len(request_ids), + seconds=70, + ) + return tuple(float(str(row["spend"])) for row in rows) + + +def _proxy_budget_spend(proxy: Gateway) -> float | None: + rows: Final = proxy.rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = %s', (LITELLM_PROXY_BUDGET_NAME,)) + return float(str(rows[0]["spend"])) if rows else None + + +def _counter_key(redis: Redis) -> str: + found: Final = tuple(str(key, "utf-8") for key in redis.scan_iter(match=f"*{GLOBAL_PROXY_SPEND_CACHE_KEY}")) + assert len(found) == 1, found + return found[0] + + +def _upstream_saw(proxy: Gateway, marker: str) -> bool: + with httpx.Client(base_url=proxy.upstream_url, timeout=5, trust_env=False) as upstream: + response: Final = upstream.get("/__observations") + response.raise_for_status() + return marker in json.dumps(object_value(response.json())["requests"]) + async def _add_request_cost(cache: UserApiKeyCache, response_cost: float) -> None: await update_cache( @@ -56,3 +153,63 @@ async def test_global_proxy_spend_adds_every_instance_cost_and_stays_cleared_aft assert await redis_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None finally: await redis_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_a_key_with_no_user_spends_toward_the_proxy_max_budget(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy(gateway, tmp_path, {}, config=_proxy_budget_config(gateway, tmp_path, max_budget=0.01)) as owned, + owned.scenario() as scenario, + ): + key: Final = scenario.key(models=[BUDGET_MODEL]) + (spent,) = _logged_spend(owned, (_served_id(owned, key),)) + assert spent > 0.01, spent + eventually( + lambda: _proxy_budget_spend(owned), + lambda recorded: recorded is not None and recorded > 0.01, + seconds=70, + ) + + marker: Final = f"over the proxy budget {uuid.uuid4().hex}" + refused: Final = _chat(owned, key, marker) + + assert refused.status_code != 200, refused.text + assert object_value(refused.json()["error"])["type"] == "budget_exceeded", refused.text + assert not _upstream_saw(owned, marker), refused.text + + +@pytest.mark.timeout(2 * OWNED_PROXY_SECONDS) +def test_two_proxies_add_every_request_to_one_proxy_spend_counter(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _proxy_budget_config(gateway, tmp_path, max_budget=1000.0) + with owned_redis(tmp_path / "redis") as cache, Redis(host=cache.host, port=cache.port) as redis: + shared_redis: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)} + with owned_proxy(gateway, tmp_path / "first", shared_redis, config=config) as first: + assert first.database_url is not None + with ( + owned_proxy( + gateway, tmp_path / "second", {**shared_redis, "DATABASE_URL": first.database_url}, config=config + ) as second, + first.scenario() as scenario, + ): + user_key: Final = scenario.key(user_id=scenario.user(), models=[BUDGET_MODEL]) + userless_key: Final = scenario.key(models=[BUDGET_MODEL]) + counter: Final = _counter_key(redis) + ttl_before: Final = redis.ttl(counter) + + request_ids: Final = tuple( + _served_id(proxy, key) for key, proxy in product((user_key, userless_key), (first, second)) + ) + total: Final = sum(_logged_spend(first, request_ids)) + persisted: Final = eventually( + lambda: _proxy_budget_spend(first), + lambda spend: spend is not None and spend == pytest.approx(total), + seconds=70, + ) + + assert persisted == pytest.approx(total) + eventually( + lambda: redis.get(counter), + lambda shared: shared is not None and float(str(shared, "utf-8")) == pytest.approx(total), + seconds=70, + ) + assert 0 < redis.ttl(counter) <= ttl_before From 07d314beb598ae4cb5d263fda3e1ecc531633112 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:05:08 -0700 Subject: [PATCH 8/9] test(responses): audit #47445 (#48902) * test(responses): audit #47445 Integration cells for billing every round of a streamed MCP auto-execute on one spend row, on a native Responses upstream and the chat bridge, across tool rounds, a rejected or broken follow-up, a post-call guardrail block, a client hang-up, the response cache, and the next bridged turn * test(responses): check a broken follow-up bills its own nonzero usage --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../integration/mcp/test_mcp_llm_endpoints.py | 365 +++++++++++++++++- 1 file changed, 351 insertions(+), 14 deletions(-) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 394221708bf..f751d1f8f4c 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from hashlib import sha256 -from itertools import chain +from itertools import chain, count from pathlib import Path from typing import Final, Literal, TypeVar @@ -50,6 +50,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" + (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE) config.write_text( yaml.safe_dump( { @@ -57,8 +58,17 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat "general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True}, "litellm_settings": { **object_value(stock["litellm_settings"]), - "cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}}, + "cost_margin_config": { + "hosted_vllm": {"fixed_amount": FIXED_MARGIN}, + "deepseek": {"fixed_amount": FIXED_MARGIN}, + }, }, + "guardrails": [ + { + "guardrail_name": STREAM_BLOCK, + "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"}, + } + ], } ) ) @@ -79,7 +89,7 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway yield private -Surface = Literal["chat", "responses", "messages", "messages_bridge"] +Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} ANSWER: Final = "the sum is 5" @@ -88,6 +98,19 @@ AUTO: Final = {**GATEWAY_REF, "require_approval": "never"} OUTAGE: Final = "bridge-outage" DIRECT: Final = "answer-without-tools" FIXED_MARGIN: Final = 0.25 +STREAM_BLOCK: Final = "mcp-stream-block" +STREAM_BLOCK_MODULE: Final = "mcp_stream_block" +STREAM_BLOCK_DETAIL: Final = "blocked by output policy" +STREAM_BLOCK_SOURCE: Final = ( + "from fastapi import HTTPException\n\n" + "from litellm.integrations.custom_guardrail import CustomGuardrail\n\n\n" + "class BlockGuard(CustomGuardrail):\n" + " async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):\n" + " async for _ in response:\n" + " pass\n" + f' raise HTTPException(status_code=400, detail="{STREAM_BLOCK_DETAIL}")\n' + " yield\n" +) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -128,6 +151,8 @@ class Billing: first: tuple[int, int] = (10, 5) follow_up: tuple[int, int] = (10, 5) follow_up_fails: bool = False + follow_up_breaks: bool = False + tool_rounds: int = 1 EVEN_ROUNDS: Final = Billing() @@ -138,6 +163,10 @@ def _fixed_turn(tool: str) -> Callable[[Mapping[str, JsonValue]], Turn]: return lambda _: Turn(tool, json.dumps(ADD), ANSWER) +def _slow_turn(tool: str, seconds: float) -> Callable[[Mapping[str, JsonValue]], Turn]: + return lambda _: Turn(tool, json.dumps({"seconds": seconds}), ANSWER) + + def _echoing_turn(body: Mapping[str, JsonValue]) -> Turn: names: Final = _tool_names(body) return Turn(names[0] if names else "", json.dumps({"query": _prompt(body)}), _tool_result_text(body) or "") @@ -183,6 +212,50 @@ def _item_events(index: int, item: Mapping[str, JsonValue]) -> tuple[dict[str, J ) +def _event_stream(events: Sequence[Mapping[str, JsonValue]]) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=tuple( + f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode() + for number, event in enumerate(events) + ), + ) + + +def _broken_follow_up(target: str, identity: str, tokens: tuple[int, int]) -> Reply: + outage: Final = {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None} + if target.endswith("/chat/completions"): + delta: Final = {"role": "assistant", "content": "the sum"} + started: Final = { + "id": f"chatcmpl-{identity}", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + } + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in (started, {"error": outage})), + ) + input_tokens, output_tokens = tokens + failed: Final[dict[str, JsonValue]] = { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1, + "status": "failed", + "model": "gpt-4o-mini", + "output": [], + "error": {"code": "server_error", "message": FOLLOW_UP_OUTAGE}, + "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": sum(tokens)}, + } + return _event_stream( + ( + {"type": "response.created", "response": {**failed, "status": "in_progress", "error": None, "usage": None}}, + {"type": "response.failed", "response": failed}, + ) + ) + + def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping[str, JsonValue]]) -> Reply: item_events: Final = tuple(chain.from_iterable(_item_events(index, item) for index, item in enumerate(items))) events: Final = ( @@ -191,13 +264,7 @@ def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping *item_events, {"type": "response.completed", "response": response}, ) - return Reply( - content_type="text/event-stream", - chunks=tuple( - f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode() - for number, event in enumerate(events) - ), - ) + return _event_stream(events) def _block_ids(message: JsonValue, kind: str, field: str) -> frozenset[str]: @@ -252,6 +319,8 @@ def _anthropic_tool_result_rule(body: Mapping[str, JsonValue]) -> Reply | None: def _model_double( plan: Callable[[Mapping[str, JsonValue]], Turn], billing: Billing = EVEN_ROUNDS, parallel: int = 1 ) -> Callable[[Request], Reply]: + served: Final = count() + def respond(request: Request) -> Reply: if request.method == "GET" and request.target.endswith("/models"): return _json({"object": "list", "data": []}) @@ -260,11 +329,12 @@ def _model_double( outage: Final = {"error": {"message": "scripted provider outage", "type": "server_error", "code": None}} return Reply(status=500, body=json.dumps(outage).encode()) turn: Final = plan(body) + call_number: Final = next(served) answered: Final = _has_tool_result(body) if answered and billing.follow_up_fails: failure: Final = {"error": {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}} return Reply(status=500, body=json.dumps(failure).encode()) - done: Final = answered or DIRECT in _prompt(body) + done: Final = DIRECT in _prompt(body) or (answered and call_number >= billing.tool_rounds) calls: Final = range(1, parallel + 1) prompt_tokens, completion_tokens = billing.follow_up if answered else billing.first identity: Final = uuid.uuid4().hex[:12] @@ -273,9 +343,15 @@ def _model_double( "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens, } + if answered and billing.follow_up_breaks: + return _broken_follow_up(request.target, identity, (prompt_tokens, completion_tokens)) if request.target.endswith("/chat/completions"): requested: Final = tuple( - {"id": f"call_{call}", "type": "function", "function": {"name": turn.tool, "arguments": turn.arguments}} + { + "id": f"call_{identity}_{call}", + "type": "function", + "function": {"name": turn.tool, "arguments": turn.arguments}, + } for call in calls ) message: Final = ( @@ -460,6 +536,7 @@ def _upstream_model(surface: Surface) -> str: "responses": "openai/responses/gpt-4o-mini", "messages": "anthropic/claude-sonnet-4-5", "messages_bridge": "hosted_vllm/gpt-4o-mini", + "responses_bridge": "deepseek/gpt-4o-mini", }[surface] @@ -475,11 +552,15 @@ def _rig( billing: Billing = EVEN_ROUNDS, prices: Mapping[str, JsonValue] | None = None, parallel: int = 1, + slow_tool_seconds: float | None = None, ) -> Iterator[Rig]: alias: Final = "llm" + uuid.uuid4().hex[:8] + plan: Final = ( + _fixed_turn(f"{alias}-add") if slow_tool_seconds is None else _slow_turn(f"{alias}-slow", slow_tool_seconds) + ) with ( - mcp_peer() as peer, - wire_server(_model_double(_fixed_turn(f"{alias}-add"), billing, parallel)) as wire, + mcp_peer(rich=slow_tool_seconds is not None) as peer, + wire_server(_model_double(plan, billing, parallel)) as wire, gateway.scenario() as scenario, ): server_id: Final = register_mcp(scenario, peer, alias) @@ -1048,6 +1129,262 @@ def test_bridged_next_turn_replays_the_mcp_turn_answer(gateway: Gateway) -> None assert ANSWER in _assistant_texts(opening), opening +StreamedUpstream = Literal["messages_bridge", "responses_bridge"] +STREAMED_UPSTREAMS: Final[tuple[StreamedUpstream, ...]] = ("messages_bridge", "responses_bridge") +STREAMED_IDS: Final = ("native", "bridged") +SLOW_TOOL_SECONDS: Final = 5.0 + + +def _events(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _completed(response: httpx.Response) -> dict[str, JsonValue]: + finished: Final = tuple( + object_value(event["response"]) for event in _events(response) if event["type"] == "response.completed" + ) + assert len(finished) == 1, response.text + return finished[0] + + +def _model_rows(call_id: str, settled: Callable[[list[dict[str, JsonValue]]], bool]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + "SELECT request_id, status, spend, prompt_tokens, completion_tokens " + 'FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type NOT IN (%s, %s)', + (call_id, "call_mcp_tool", "list_mcp_tools"), + ), + settled, + seconds=70, + return_last_on_timeout=True, + ) + + +def _billed(rows: Sequence[Mapping[str, JsonValue]]) -> float: + return sum(float(str(row["spend"])) for row in rows) + + +def _bills_at_least(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]: + return lambda rows: bool(rows) and _billed(rows) >= expected - 1e-9 + + +def _statuses(rows: Sequence[Mapping[str, JsonValue]]) -> list[str]: + return sorted(str(row["status"]) for row in rows) + + +def _logs_a_failure_billing(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]: + return lambda rows: "failure" in _statuses(rows) and _billed(rows) >= expected - 1e-9 + + +def _limited_key(rig: Rig) -> str: + return rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]}, max_parallel_requests=1) + + +def _answers_directly(rig: Rig, key: str) -> int: + return rig.send(key, [AUTO], prompt=f"{DIRECT} {rig.alias}").status_code + + +STREAMED_ROUNDS: Final = ( + pytest.param("messages_bridge", 1, id="native"), + pytest.param("responses_bridge", 1, id="bridged"), + pytest.param("messages_bridge", 2, id="native-two-tool-rounds"), + pytest.param("responses_bridge", 2, id="bridged-two-tool-rounds"), + pytest.param("messages_bridge", 0, id="native-no-tool-call"), + pytest.param("responses_bridge", 0, id="bridged-no-tool-call"), +) + + +@pytest.mark.parametrize(("upstream", "tool_rounds"), STREAMED_ROUNDS) +def test_streaming_responses_auto_execute_bills_every_round_once_on_one_success_row( + gateway: Gateway, upstream: StreamedUpstream, tool_rounds: int +) -> None: + rounds: Final = (FIRST_ROUND, *(FOLLOW_UP_ROUND,) * tool_rounds) + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, tool_rounds=max(tool_rounds, 1)) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + prompt: Final = None if tool_rounds else f"{DIRECT} {rig.alias}" + response: Final = rig.send(_granted_key(rig), [AUTO], prompt=prompt, stream=True) + assert response.status_code == 200, response.text + completed: Final = _completed(response) + assert rig.final_text(completed) == ANSWER, response.text + assert len(rig.upstream_tools()) == len(rounds) + assert len(_peer_add_calls(rig.peer)) == tool_rounds + expected: Final = _price(*rounds) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected)) + assert [row["status"] for row in rows] == ["success"], rows + row: Final = rows[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + sum(prompt_tokens for prompt_tokens, _ in rounds), + sum(completion_tokens for _, completion_tokens in rounds), + ), row + usage: Final = object_value(completed["usage"]) + assert (usage["input_tokens"], usage["output_tokens"]) == (row["prompt_tokens"], row["completion_tokens"]) + assert float(str(row["spend"])) == pytest.approx(expected), row + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_rejected_follow_up_bills_round_one_once( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True) + assert FOLLOW_UP_OUTAGE in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + expected: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected)) + assert _billed(rows) == pytest.approx(expected), rows + + +BROKEN_FOLLOW_UPS: Final = ( + pytest.param("messages_bridge", FOLLOW_UP_ROUND, id="native"), + pytest.param("responses_bridge", None, id="bridged"), +) + + +@pytest.mark.parametrize(("upstream", "reported_usage"), BROKEN_FOLLOW_UPS) +def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_and_itself_on_a_failure_row( + gateway: Gateway, upstream: StreamedUpstream, reported_usage: tuple[int, int] | None +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True) + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + failed_round: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(failed_round) > 0, rows + assert reported_usage in (None, failed_round), rows + assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + key: Final = _granted_key(rig) + one_round: Final = rig.send(key, [], prompt=f"{DIRECT} {rig.alias}", stream=True, guardrails=[STREAM_BLOCK]) + assert STREAM_BLOCK_DETAIL in one_round.text, one_round.text + one_round_rows: Final = _model_rows( + one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN) + ) + rig.wire.drain() + response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK]) + assert STREAM_BLOCK_DETAIL in response.text, response.text + assert len(rig.upstream_tools()) == 2 and len(_peer_add_calls(rig.peer)) == 1 + expected: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(expected)) + assert _statuses(rows) == _statuses(one_round_rows), (rows, one_round_rows) + assert _billed(rows) == pytest.approx(expected), rows + + +@pytest.mark.parametrize( + ("surface", "stream"), + (("responses", True), ("responses", False), ("chat", False)), + ids=("responses-stream", "responses", "chat"), +) +def test_auto_executed_tool_call_runs_inside_the_requests_one_parallel_slot( + gateway: Gateway, surface: Surface, stream: bool +) -> None: + with _rig(gateway, surface, slow_tool_seconds=SLOW_TOOL_SECONDS) as rig: + key: Final = _limited_key(rig) + with ThreadPoolExecutor(max_workers=1) as pool: + running: Final = pool.submit(rig.send, key, [AUTO], None, stream=stream) + eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30) + during: Final = _answers_directly(rig, key) + served: Final = running.result(timeout=90) + assert during == 429 + assert served.status_code == 200 and ANSWER in served.text, served.text + eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=30) + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_hung_up_during_the_tool_call_bills_each_served_round_once_and_frees_the_slot( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig( + gateway, + "responses", + upstream=upstream, + billing=billing, + prices=PRICES, + slow_tool_seconds=SLOW_TOOL_SECONDS, + ) as rig: + key: Final = _limited_key(rig) + with rig.gateway.client.stream( + "POST", + "/v1/responses", + headers={"Authorization": f"Bearer {key}"}, + json={"model": rig.model, "input": f"add {rig.alias}", "tools": [AUTO], "stream": True}, + timeout=90, + ) as hung_up: + assert hung_up.status_code == 200 + call_id: Final = hung_up.headers["x-litellm-call-id"] + eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30) + eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=60) + posted: Final = tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + served: Final = tuple( + FOLLOW_UP_ROUND if _has_tool_result(body) else FIRST_ROUND for body in posted if DIRECT not in _prompt(body) + ) + expected: Final = _price(*served) + FIXED_MARGIN + rows: Final = _model_rows(call_id, _bills_at_least(expected)) + assert served and rows, (served, rows) + assert _billed(rows) == pytest.approx(expected), (served, rows) + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_repeated_from_the_response_cache_bills_nothing( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + key: Final = _granted_key(rig) + paid: Final = rig.send(key, [AUTO], stream=True) + assert paid.status_code == 200 and rig.final_text(_completed(paid)) == ANSWER, paid.text + paid_price: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN + _model_rows(paid.headers["x-litellm-call-id"], _bills_at_least(paid_price)) + assert len(rig.upstream_tools()) == 2 + cached: Final = rig.send(key, [AUTO], stream=True) + assert cached.status_code == 200 and rig.final_text(_completed(cached)) == ANSWER, cached.text + assert rig.upstream_tools() == () + rows: Final = _model_rows( + cached.headers["x-litellm-call-id"], lambda found: "success" in {row["status"] for row in found} + ) + assert "success" in {row["status"] for row in rows} and _billed(rows) == 0.0, rows + + +def test_bridged_next_turn_after_a_streamed_mcp_turn_replays_its_answer(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", upstream="responses_bridge", billing=billing, prices=PRICES) as rig: + key: Final = _granted_key(rig) + first: Final = rig.send(key, [AUTO], stream=True) + assert first.status_code == 200, first.text + first_id: Final = string_value(_completed(first)["id"]) + _model_rows(first.headers["x-litellm-call-id"], _bills_at_least(_price(FIRST_ROUND, FOLLOW_UP_ROUND))) + rig.wire.drain() + second: Final = rig.send(key, [AUTO], previous_response_id=first_id, cache={"no-cache": True}) + assert second.status_code == 200, second.text + opening: Final = next( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert ANSWER in _assistant_texts(opening), opening + + Bridge = Literal["chat", "responses", "messages"] BRIDGES: Final[tuple[Bridge, ...]] = ("chat", "responses") Client = Literal["sync", "async"] From 3b9a59348f564d7746cb05a9e16341f42321a219 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:06:58 -0700 Subject: [PATCH 9/9] fix(vertex_ai)!: send the caller's anthropic-version in the Vertex Claude /v1/messages body, not the header (#48814) * fix(vertex_ai): drop the caller's anthropic-version header on Claude /v1/messages Vertex reads the API version from the body's anthropic_version and answers 400 when a forwarded anthropic-version header names a different one, which anthropic SDK 2.x always does. Strip the header on the Python and LITELLM_RUST Vertex messages routes. * fix(vertex_ai): carry the caller's anthropic-version into the Claude /v1/messages body Vertex reads the API version from the body alone, so the caller's anthropic-version header now becomes the body's anthropic_version instead of a pinned vertex-2023-10-16. The 2023-01-01 and 2023-06-01 first-party versions keep vertex-2023-10-16, and the header itself is still never sent. Same change on the Python proxy and the Rust gateway. * fix(vertex_ai): follow drop_params for an anthropic-version Vertex does not accept Only versions Vertex accepts in the body (vertex-2023-10-16, 2026-10-21) are carried into anthropic_version. 2023-06-01, 2023-01-01, a blank header, or no header keep vertex-2023-10-16. Any other value drops to vertex-2023-10-16 under drop_params and answers a 400 naming drop_params without it, on the Python proxy and the Rust gateway. * refactor(vertex_ai): return the unsupported anthropic-version error as a value and keep Rust helpers and tests in place * fix(rust gateway): resolve os.environ drop_params and accept the same flag words as Python The gateway read drop_params: os.environ/NAME and words like yes as off, so the drop_params opt-out from the new anthropic-version 400 did not work there * docs(rust router): name Router::from_config in the crate README --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + .../crates/core-utils/src/serde_compat.rs | 41 ++++- litellm-rust/crates/gateway/src/lib.rs | 8 +- .../crates/inference-messages/src/handler.rs | 88 ++++++++- .../src/base_llm/messages/transformation.rs | 9 + .../llms/src/vertex_ai/messages/AGENTS.md | 3 +- .../src/vertex_ai/messages/transformation.rs | 173 ++++++++++++++++-- litellm-rust/crates/router/Cargo.toml | 1 + litellm-rust/crates/router/README.md | 2 +- litellm-rust/crates/router/src/lib.rs | 88 ++++++++- litellm-rust/crates/router/tests/router.rs | 8 +- .../anthropic/messages/transformation.py | 59 +++++- ...artner_models_anthropic_messages_config.py | 83 ++++++++- 13 files changed, 522 insertions(+), 42 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 10e590b5006..66a345fc438 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4387,6 +4387,7 @@ name = "litellm-router" version = "0.1.0" dependencies = [ "litellm-config", + "litellm-core-utils", "litellm-inference", "litellm-inference-messages", "litellm-router-types", diff --git a/litellm-rust/crates/core-utils/src/serde_compat.rs b/litellm-rust/crates/core-utils/src/serde_compat.rs index 3e4d82d3a3e..6f40a4af60a 100644 --- a/litellm-rust/crates/core-utils/src/serde_compat.rs +++ b/litellm-rust/crates/core-utils/src/serde_compat.rs @@ -1,13 +1,25 @@ -pub fn parse_str_bool(value: &str) -> Option { - let token = value.trim_matches(|character: char| { +fn python_strip(value: &str) -> &str { + value.trim_matches(|character: char| { character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}') - }); + }) +} + +pub fn parse_str_bool(value: &str) -> Option { + let token = python_strip(value); if token.eq_ignore_ascii_case("true") { return Some(true); } token.eq_ignore_ascii_case("false").then_some(false) } +pub fn parse_pydantic_bool(value: &str) -> Option { + match python_strip(value).to_ascii_lowercase().as_str() { + "1" | "on" | "t" | "true" | "y" | "yes" => Some(true), + "0" | "off" | "f" | "false" | "n" | "no" => Some(false), + _ => None, + } +} + /// `redis-py` string Booleans: only `1`, `true`, and `yes` (case-insensitive) are true. pub fn parse_redis_bool(value: &str) -> bool { value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes") @@ -34,4 +46,27 @@ mod tests { ) { assert_eq!(parse_str_bool(input), expected, "{input:?}"); } + + #[rstest] + #[case::yes(" YES ", Some(true))] + #[case::on("On", Some(true))] + #[case::one("1", Some(true))] + #[case::t("t", Some(true))] + #[case::y("y", Some(true))] + #[case::control_whitespace_true("\u{1c}TRUE\u{1f}", Some(true))] + #[case::no("No", Some(false))] + #[case::off("OFF", Some(false))] + #[case::zero("0", Some(false))] + #[case::f("F", Some(false))] + #[case::n("n", Some(false))] + #[case::false_word("false", Some(false))] + #[case::two("2", None)] + #[case::empty("", None)] + #[case::unknown("maybe", None)] + fn boolean_words_follow_pydantic_lax_bool_parsing( + #[case] input: &str, + #[case] expected: Option, + ) { + assert_eq!(parse_pydantic_bool(input), expected, "{input:?}"); + } } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index b8ff8e46033..12a9bf9110a 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -37,6 +37,7 @@ pub fn build_inference(config: &Config) -> Result, Error> { .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) }; + let router = ModelRouter::from_config(config, &lookup); let settings = &config.litellm_settings; let http_settings = HttpSettings::from_layers([ HttpSettingsLayer::from_environment(&lookup), @@ -64,12 +65,7 @@ pub fn build_inference(config: &Config) -> Result, Error> { Arc::new(EnvironmentSecrets::python_compatible(client)), )); let resources = CoreResources::new(pool); - Ok(Arc::new(Gateway::new( - resources, - http, - secrets, - ModelRouter::from_model_list(&config.model_list), - )?)) + Ok(Arc::new(Gateway::new(resources, http, secrets, router)?)) } pub async fn build_redis(config: &Config) -> Result { diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index fa802d20587..9d0bf09197e 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -141,9 +141,13 @@ async fn provider_call( &env_lookup, )?; - let body = config.transform_anthropic_messages_request( - trimmed, - &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), + let body = config.request_body( + config.transform_anthropic_messages_request( + trimmed, + &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), + )?, + &validated.headers, + shaping.settings.drop_params, )?; let environment = ValidatedEnvironment { headers: config.request_headers( @@ -514,6 +518,84 @@ mod tests { ); } + fn vertex_call_requesting( + shaping: MessagesShaping, + requested: &str, + drop_params: bool, + ) -> MessagesCall { + MessagesCall { + body: body( + json!({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + ), + api_key: None, + api_base: Some("https://vertex.test".into()), + custom_llm_provider: Some("vertex_ai".into()), + litellm_params: Default::default(), + extra_headers: Some(Map::from_iter([ + ("authorization".to_string(), json!("Bearer vertex-token")), + ("anthropic-version".to_string(), json!(requested)), + ])), + provider_specific_header: None, + timeout: None, + shaping: MessagesShaping { + settings: MessagesSettings { + drop_params, + ..shaping.settings + }, + ..shaping + }, + } + } + + fn vertex_project(name: &str) -> Option { + (name == "VERTEXAI_PROJECT").then(|| "version-project".to_string()) + } + + #[rstest] + #[case::sdk_2("2026-10-21", false, "2026-10-21")] + #[case::sdk_1("2023-06-01", false, "vertex-2023-10-16")] + #[case::unknown_under_drop_params("latest", true, "vertex-2023-10-16")] + #[tokio::test] + async fn vertex_gets_the_requested_version_in_the_body_alone( + shaping: MessagesShaping, + #[case] requested: &str, + #[case] drop_params: bool, + #[case] body_version: &str, + ) { + let call = provider_call_with_secrets( + vertex_call_requesting(shaping, requested, drop_params), + &vertex_project, + ) + .await + .unwrap(); + let version_headers: Vec<&str> = call + .wire + .headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("anthropic-version")) + .map(|(_, value)| value.as_str()) + .collect(); + assert_eq!( + (version_headers, &call.wire.body["anthropic_version"]), + (Vec::<&str>::new(), &json!(body_version)) + ); + } + + #[rstest] + #[tokio::test] + async fn vertex_refuses_an_unsupported_version_without_drop_params(shaping: MessagesShaping) { + assert!(matches!( + provider_call_with_secrets( + vertex_call_requesting(shaping, "latest", false), + &vertex_project + ) + .await, + Err(Error::Llm(litellm_llms::Error::InvalidRequest( + litellm_llms::ErrorDetail::UnsupportedParameter { .. } + ))) + )); + } + #[rstest] #[tokio::test] async fn missing_credentials_are_reported_before_the_request_is_transformed( diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 8a88af40d6b..744ce89afdb 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -77,6 +77,15 @@ pub trait BaseMessagesConfig: Sync { headers } + fn request_body( + &self, + request: MessagesRequest, + _headers: &Headers, + _drop_params: bool, + ) -> Result { + Ok(request) + } + /// The JSON that goes on the wire, for a host whose body differs from the typed request: /// Python's configs `pop("model")` when the model is addressed by the URL. fn wire_body(&self, body: Value) -> Value { diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md b/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md index b5d637d6bc7..d95da65ac9d 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md @@ -1,6 +1,7 @@ # rules -- Claude on Vertex AI: reuses `anthropic/messages` shaping, moves the model into the publisher URL, puts the Vertex `anthropic_version` in the body, and authenticates with a Google access token +- Claude on Vertex AI: reuses `anthropic/messages` shaping, moves the model into the publisher URL, and authenticates with a Google access token +- Vertex reads the API version only from the body's `anthropic_version` and answers 400 when an `anthropic-version` header names a different one, so the caller's header never reaches Vertex. A version Vertex accepts in the body is carried there, `2023-06-01`, `2023-01-01`, or no header get `vertex-2023-10-16`, and any other version drops under `drop_params` and is refused without it - Applies Vertex's beta policy and strips `output_config.effort` for models that do not accept it # references diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 2e77dbc1ca8..46e5a68dcb8 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -3,10 +3,10 @@ use litellm_auth_gcp::{ VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, secret_names, }; -use litellm_http::request::has_header; +use litellm_http::request::{has_header, header_value, without_headers}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool, OutputConfig}, - providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet}, + providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, VERSION_HEADER}, recognized::Recognized, }; use litellm_router_types::LitellmParams; @@ -36,6 +36,8 @@ use litellm_llms_types::providers::vertex_ai::{ }; const WEB_SEARCH_TOOL_PREFIX: &str = "web_search"; +const VERSIONS_SERVED_AS_VERTEX_DEFAULT: [&str; 3] = ["", "2023-01-01", "2023-06-01"]; +const VERSIONS_VERTEX_ACCEPTS_IN_BODY: [&str; 2] = [ANTHROPIC_VERSION, "2026-10-21"]; /// Claude on Vertex AI, Python's `VertexAIPartnerModelsAnthropicMessagesConfig`: the Anthropic /// payload addressed to a project and location, authenticated with a Google access token. @@ -90,12 +92,6 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { .params .output_config .and_then(|config| sanitize_output_config(config, accepts_effort)), - extra: request - .params - .extra - .into_iter() - .chain([("anthropic_version".to_string(), json!(ANTHROPIC_VERSION))]) - .collect(), ..request.params }, ..request @@ -133,7 +129,35 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - merge_beta_headers(headers, vertex_feature_betas(request)) + merge_beta_headers( + without_headers(headers, &[VERSION_HEADER]), + vertex_feature_betas(request), + ) + } + + fn request_body( + &self, + request: MessagesRequest, + headers: &Headers, + drop_params: bool, + ) -> Result { + let version = vertex_version( + &request.model, + header_value(headers, VERSION_HEADER), + drop_params, + )?; + Ok(MessagesRequest { + params: MessagesOptionalParams { + extra: request + .params + .extra + .into_iter() + .chain([("anthropic_version".to_string(), json!(version))]) + .collect(), + ..request.params + }, + ..request + }) } /// Vertex addresses the model in the URL and rejects it in the body. @@ -225,6 +249,26 @@ fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet { ) } +fn vertex_version<'a>( + model: &str, + requested: Option<&'a str>, + drop_params: bool, +) -> Result<&'a str, Error> { + match requested.map(str::trim).unwrap_or_default() { + version if VERSIONS_SERVED_AS_VERTEX_DEFAULT.contains(&version) => Ok(ANTHROPIC_VERSION), + version if VERSIONS_VERTEX_ACCEPTS_IN_BODY.contains(&version) => Ok(version), + _ if drop_params => Ok(ANTHROPIC_VERSION), + version => Err(Error::InvalidRequest( + crate::ErrorDetail::UnsupportedParameter { + model: model.into(), + param: VERSION_HEADER.into(), + value: version.into(), + hint: String::new(), + }, + )), + } +} + #[cfg(test)] mod tests { use rstest::rstest; @@ -471,6 +515,28 @@ mod tests { ); } + #[rstest] + #[case::sdk_spelling("anthropic-version")] + #[case::any_case("Anthropic-Version")] + fn a_forwarded_anthropic_version_never_reaches_vertex(#[case] name: &str) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + assert_eq!( + VERTEX_ANTHROPIC_MESSAGES_CONFIG.request_headers( + vec![ + (name.to_string(), "2026-10-21".to_string()), + ("x-request-id".to_string(), "req-1".to_string()), + ], + &request, + ), + vec![("x-request-id".to_string(), "req-1".to_string())] + ); + } + fn transformed(value: Value, context: &MessagesTransformContext) -> Value { let request: MessagesRequest = serde_json::from_value(value).unwrap(); serde_json::to_value( @@ -482,16 +548,87 @@ mod tests { } #[rstest] - fn the_body_carries_the_vertex_version_and_the_wire_drops_the_model() { - let body = transformed( - json!({ - "model": "claude-sonnet-4-5@20250929", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hi"}] - }), - &MessagesTransformContext::default(), + #[case::sdk_2(&[("anthropic-version", "2026-10-21")], false, "2026-10-21")] + #[case::any_case(&[("Anthropic-Version", "2026-10-21")], false, "2026-10-21")] + #[case::vertex_name(&[("anthropic-version", "vertex-2023-10-16")], false, "vertex-2023-10-16")] + #[case::sdk_1(&[("anthropic-version", "2023-06-01")], false, "vertex-2023-10-16")] + #[case::first_version(&[("anthropic-version", "2023-01-01")], false, "vertex-2023-10-16")] + #[case::blank(&[("anthropic-version", " ")], false, "vertex-2023-10-16")] + #[case::no_header(&[], false, "vertex-2023-10-16")] + #[case::unknown_under_drop_params(&[("anthropic-version", "latest")], true, "vertex-2023-10-16")] + fn the_requested_version_goes_in_the_body( + #[case] forwarded: &[(&str, &str)], + #[case] drop_params: bool, + #[case] body_version: &str, + ) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + let headers: Headers = forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + let body = serde_json::to_value( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_body(request, &headers, drop_params) + .unwrap(), + ) + .unwrap(); + assert_eq!(body["anthropic_version"], json!(body_version)); + } + + #[rstest] + #[case::unreleased("2025-01-01")] + #[case::alias("latest")] + fn an_unsupported_version_is_refused_without_drop_params(#[case] requested: &str) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + let headers: Headers = vec![("anthropic-version".to_string(), requested.to_string())]; + let Error::InvalidRequest(detail) = VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_body(request, &headers, false) + .unwrap_err() + else { + panic!("expected an invalid request"); + }; + assert_eq!( + detail.to_string(), + format!( + "claude-sonnet-4-5 does not support anthropic-version={requested}. \ + To drop unsupported params, set `litellm.drop_params = True`." + ) ); - assert_eq!(body["anthropic_version"], json!("vertex-2023-10-16")); + } + + #[rstest] + fn the_body_carries_the_vertex_version_and_the_wire_drops_the_model() { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5@20250929", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + let body = serde_json::to_value( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_body( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request( + request, + &MessagesTransformContext::default(), + ) + .unwrap(), + &Vec::new(), + false, + ) + .unwrap(), + ) + .unwrap(); assert_eq!(body["model"], json!("claude-sonnet-4-5@20250929")); let wire = VERTEX_ANTHROPIC_MESSAGES_CONFIG.wire_body(body); assert_eq!(wire.get("model"), None); diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index a356bedcf16..988d1b68805 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-config.workspace = true +litellm-core-utils.workspace = true litellm-inference.workspace = true litellm-inference-messages.workspace = true litellm-router-types.workspace = true diff --git a/litellm-rust/crates/router/README.md b/litellm-rust/crates/router/README.md index 972d37d14be..a734e634104 100644 --- a/litellm-rust/crates/router/README.md +++ b/litellm-rust/crates/router/README.md @@ -1,4 +1,4 @@ -`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_model_list(&config.model_list)` maps configured public names to provider deployments. Programmatic callers can collect `(String, Deployment)` entries into a `Router` +`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_config(&config, &env)` maps configured public names to provider deployments and reads `drop_params` from `litellm_settings` and each deployment, resolving an `os.environ/NAME` value through `env`. Programmatic callers can collect `(String, Deployment)` entries into a `Router` Lookup is exact and returns `None` for an unknown name. This extraction preserves the gateway's existing behavior: the last entry wins when public names repeat. Multiple deployments per model group, routing strategies, retries, cooldowns, and fallbacks are not implemented yet diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index fb48213e8a7..b75c1a14446 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -2,16 +2,33 @@ mod deployment; use std::collections::HashMap; -use litellm_config::Model; +use litellm_config::{Config, Spelled}; +use litellm_core_utils::{serde_compat::parse_pydantic_bool, settings::Lookup}; +use litellm_inference_messages::{MessagesSettings, MessagesShaping}; pub use deployment::Deployment; #[derive(Clone, Debug, Default)] pub struct Router(HashMap); +fn is_on(flag: Option<&Spelled>, env: &impl Lookup) -> bool { + match flag { + Some(Spelled::Value(on)) => *on, + Some(Spelled::Text(text)) => match text.strip_prefix("os.environ/") { + Some(name) => env + .get(name) + .is_some_and(|value| parse_pydantic_bool(&value) == Some(true)), + None => parse_pydantic_bool(text) == Some(true), + }, + None => false, + } +} + impl Router { - pub fn from_model_list(model_list: &[Model]) -> Self { - model_list + pub fn from_config(config: &Config, env: &impl Lookup) -> Self { + let drops_params_globally = is_on(config.litellm_settings.drop_params.as_ref(), env); + config + .model_list .iter() .map(|model| { ( @@ -26,6 +43,14 @@ impl Router { api_base: model.litellm_params.api_base.clone(), custom_llm_provider: model.litellm_params.custom_llm_provider.clone(), litellm_params: model.litellm_params.clone(), + shaping: MessagesShaping { + settings: MessagesSettings { + drop_params: drops_params_globally + || is_on(model.litellm_params.drop_params.as_ref(), env), + ..MessagesSettings::default() + }, + ..MessagesShaping::default() + }, ..Deployment::default() }, ) @@ -43,3 +68,60 @@ impl FromIterator<(String, Deployment)> for Router { Self(entries.into_iter().collect()) } } + +#[cfg(test)] +mod tests { + use litellm_config::Config; + use rstest::rstest; + + use super::Router; + + #[rstest] + #[case::unset("", "", false)] + #[case::deployment_flag("", "\n drop_params: true", true)] + #[case::deployment_word("", "\n drop_params: \" True \"", true)] + #[case::deployment_off("", "\n drop_params: false", false)] + #[case::deployment_unknown_word("", "\n drop_params: maybe", false)] + #[case::deployment_yes("", "\n drop_params: yes", true)] + #[case::deployment_no("", "\n drop_params: \"no\"", false)] + #[case::deployment_env_reference("", "\n drop_params: os.environ/DROP_PARAMS_FLAG", true)] + #[case::deployment_unset_env_reference( + "", + "\n drop_params: os.environ/MISSING_FLAG", + false + )] + #[case::global_env_reference( + "litellm_settings:\n drop_params: os.environ/DROP_PARAMS_FLAG\n", + "", + true + )] + #[case::global_flag("litellm_settings:\n drop_params: true\n", "", true)] + #[case::global_word("litellm_settings:\n drop_params: \"true\"\n", "", true)] + #[case::global_off("litellm_settings:\n drop_params: false\n", "", false)] + #[case::global_over_deployment_off( + "litellm_settings:\n drop_params: true\n", + "\n drop_params: false", + true + )] + fn configured_drop_params_reach_the_deployment( + #[case] settings: &str, + #[case] deployment: &str, + #[case] expected: bool, + ) { + let config = Config::from_yaml(&format!( + "{settings}model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model{deployment}" + )) + .unwrap(); + + assert_eq!( + Router::from_config(&config, &|name: &str| (name == "DROP_PARAMS_FLAG") + .then(|| " Yes ".to_owned())) + .get("public-model") + .unwrap() + .shaping + .settings + .drop_params, + expected + ); + } +} diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index bbaeb806b96..03e78ab9d11 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -16,7 +16,7 @@ fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n {parameters}" )) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_config(&config, &|_: &str| None); let deployment = router.get(&config.model_list[0].model_name).unwrap(); let params = &config.model_list[0].litellm_params; @@ -44,7 +44,7 @@ fn configuration_types_the_provider_params_of_a_deployment() { azure_ad_token: ignored-by-every-group", ) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_config(&config, &|_: &str| None); let params = &router.get("public-model").unwrap().litellm_params; assert_eq!( @@ -74,7 +74,7 @@ fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) model: provider/b", ) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_config(&config, &|_: &str| None); assert_eq!(router.get(name).map(|entry| entry.model.as_str()), expected); } @@ -84,7 +84,7 @@ fn empty_configuration_has_no_deployment() { let config = Config::from_yaml("model_list: []").unwrap(); assert!( - Router::from_model_list(&config.model_list) + Router::from_config(&config, &|_: &str| None) .get("") .is_none() ); diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index 6d8eb228786..1f0ee728cd1 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -1,6 +1,8 @@ from collections.abc import Mapping, Sequence from typing import Final +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, @@ -17,6 +19,41 @@ from litellm.types.router import GenericLiteLLMParams from ....vertex_llm_base import VertexBase from ..output_params_utils import sanitize_vertex_anthropic_output_params +_VERSION_HEADER_VERTEX_READS_FROM_BODY: Final = "anthropic-version" +_VERTEX_DEFAULT_VERSION: Final = "vertex-2023-10-16" +_VERSIONS_SERVED_AS_VERTEX_DEFAULT: Final = frozenset({"", "2023-01-01", "2023-06-01"}) +_VERSIONS_VERTEX_ACCEPTS_IN_BODY: Final = frozenset({_VERTEX_DEFAULT_VERSION, "2026-10-21"}) + + +def _requested_version(headers: Mapping[str, object]) -> str: + requested: Final = next( + (value for name, value in headers.items() if name.lower() == _VERSION_HEADER_VERTEX_READS_FROM_BODY), + "", + ) + return requested.strip() if isinstance(requested, str) else "" + + +def _vertex_version(headers: Mapping[str, object], model: str, drop_params: bool) -> str | UnsupportedParamsError: + requested: Final = _requested_version(headers) + if requested in _VERSIONS_SERVED_AS_VERTEX_DEFAULT: + return _VERTEX_DEFAULT_VERSION + if requested in _VERSIONS_VERTEX_ACCEPTS_IN_BODY: + return requested + if drop_params: + return _VERTEX_DEFAULT_VERSION + return UnsupportedParamsError( + message=( + f"{model} does not support {_VERSION_HEADER_VERTEX_READS_FROM_BODY}={requested}. " + "To drop unsupported params, set `litellm.drop_params = True`." + ), + llm_provider="vertex_ai", + model=model, + ) + + +def _without_version_header(headers: Mapping[str, str]) -> dict[str, str]: # mutable-ok: the caller sends these headers + return {name: value for name, value in headers.items() if name.lower() != _VERSION_HEADER_VERTEX_READS_FROM_BODY} + class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, VertexBase): @property @@ -198,7 +235,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert messages: list[dict], anthropic_messages_optional_request_params: dict, litellm_params: GenericLiteLLMParams, - headers: dict, + headers: dict[str, object], # mutable-ok: mirrors the provider transform contract ) -> dict: anthropic_messages_request: Final = super().transform_anthropic_messages_request( model=model, @@ -212,10 +249,28 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert self._remove_scope_from_cache_control(anthropic_messages_request) - anthropic_messages_request["anthropic_version"] = "vertex-2023-10-16" + version: Final = _vertex_version( + headers, model=model, drop_params=drops_unsupported_params(litellm_params.drop_params) + ) + if isinstance(version, UnsupportedParamsError): + raise version + anthropic_messages_request["anthropic_version"] = version anthropic_messages_request.pop("model", None) # do not pass model in request body to vertex ai sanitize_vertex_anthropic_output_params(anthropic_messages_request, model) return anthropic_messages_request + + def sign_request( + self, + headers: Mapping[str, str], + optional_params: Mapping[str, object], + request_data: Mapping[str, object], + api_base: str, + api_key: str | None = None, + model: str | None = None, + stream: bool | None = None, + fake_stream: bool | None = None, + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the caller sends these headers + return _without_version_header(headers), None diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 6c10d524531..5916ce9f7a1 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -6,12 +6,14 @@ from datetime import datetime, timedelta from queue import SimpleQueue from threading import get_ident from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import requests import responses +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) @@ -94,6 +96,85 @@ async def test_messages_auth_load_and_refresh_leave_the_event_loop(stream: bool, assert headers == {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"} +def _vertex_client() -> MagicMock: + vertex_response: Final = MagicMock() + vertex_response.json.return_value = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": "claude-sonnet-5", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=vertex_response) + return client + + +async def _send_to_vertex(client: MagicMock, caller_headers: Mapping[str, str], drop_params: bool) -> None: + with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth: + oauth.add(responses.POST, "https://oauth2.googleapis.com/token", json={"access_token": "vertex-token", "expires_in": 3600}) + await litellm.anthropic.messages.acreate( + model="vertex_ai/claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + vertex_ai_project="version-project", + vertex_ai_location="global", + vertex_credentials=json.dumps( + { + "type": "authorized_user", + "client_id": "version-header-client", + "client_secret": "synthetic-secret", + "refresh_token": "synthetic-refresh", + } + ), + client=client, + drop_params=drop_params, + provider_specific_header={ + "custom_llm_provider": "vertex_ai", + "extra_headers": {**caller_headers, "x-request-id": "req-1"}, + }, + ) + + +@pytest.mark.parametrize( + ("caller_headers", "drop_params", "body_version"), + ( + ({"anthropic-version": "2026-10-21"}, False, "2026-10-21"), + ({"Anthropic-Version": "2026-10-21"}, False, "2026-10-21"), + ({"anthropic-version": "vertex-2023-10-16"}, False, "vertex-2023-10-16"), + ({"anthropic-version": "2023-06-01"}, False, "vertex-2023-10-16"), + ({"anthropic-version": "2023-01-01"}, False, "vertex-2023-10-16"), + ({"anthropic-version": " "}, False, "vertex-2023-10-16"), + ({}, False, "vertex-2023-10-16"), + ({"anthropic-version": "latest"}, True, "vertex-2023-10-16"), + ), +) +@pytest.mark.asyncio +async def test_caller_anthropic_version_reaches_vertex_only_in_the_body( + caller_headers: dict[str, str], drop_params: bool, body_version: str +) -> None: + client: Final = _vertex_client() + await _send_to_vertex(client, caller_headers, drop_params) + + sent: Final = client.post.call_args.kwargs + sent_header_names: Final = {name.lower() for name in sent["headers"]} + assert "anthropic-version" not in sent_header_names + assert sent["headers"]["x-request-id"] == "req-1" + assert json.loads(sent["data"])["anthropic_version"] == body_version + + +@pytest.mark.parametrize("requested", ("2025-01-01", "latest")) +@pytest.mark.asyncio +async def test_unsupported_anthropic_version_is_refused_without_drop_params(requested: str) -> None: + client: Final = _vertex_client() + with pytest.raises(litellm.BadRequestError, match=f"anthropic-version={requested}. To drop unsupported params"): + await _send_to_vertex(client, {"anthropic-version": requested}, drop_params=False) + client.post.assert_not_awaited() + + def test_validate_environment_uses_vertex_ai_location(): config = VertexAIPartnerModelsAnthropicMessagesConfig() headers = {}