From f796036af031f5b340432e8f2f0fb447fafd4aee Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 16 Apr 2026 13:52:39 -0700 Subject: [PATCH 01/11] feat(proxy): add --reload flag for uvicorn hot reload (dev only) Opt-in CLI flag, off by default, no env var. Only affects the uvicorn run path; gunicorn/hypercorn paths and prod (which doesn't pass the flag) are unaffected. --- litellm/proxy/proxy_cli.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c638e294268..f1e5938c1f4 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -577,6 +577,12 @@ class ProxyInitializationHelpers: help="Exit with error if database migration fails on startup.", envvar="ENFORCE_PRISMA_MIGRATION_CHECK", ) +@click.option( + "--reload", + is_flag=True, + default=False, + help="Enable uvicorn hot reload (dev only). Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.", +) def run_server( # noqa: PLR0915 host, port, @@ -618,6 +624,7 @@ def run_server( # noqa: PLR0915 keepalive_timeout, max_requests_before_restart, enforce_prisma_migration_check: bool, + reload: bool, ): if setup: from litellm.setup_wizard import run_setup_wizard @@ -954,6 +961,9 @@ def run_server( # noqa: PLR0915 if loop_type: uvicorn_args["loop"] = loop_type + if reload: + uvicorn_args["reload"] = True + uvicorn.run( **uvicorn_args, workers=num_workers, From a19bff4ca67dcf4e721df9d831c9d073aa26acd2 Mon Sep 17 00:00:00 2001 From: nhyy244 <106547304+nhyy244@users.noreply.github.com> Date: Mon, 20 Apr 2026 23:49:41 +0200 Subject: [PATCH 02/11] Feature/add audio support for scaleway (#26110) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(scaleway): add SCALEWAY to LlmProviders enum * feat(scaleway): add audio transcription config and dispatch wiring Co-Authored-By: Claude Sonnet 4.6 * test(scaleway): add behavior tests for audio transcription config Co-Authored-By: Claude Sonnet 4.6 * chore(scaleway): advertise audio_transcriptions in endpoint-support JSON * docs(scaleway): document audio transcription support * fix(scaleway): address PR review — plain-text response_format + missing-key fail-fast Co-Authored-By: Claude Sonnet 4.6 * test(scaleway): cover new response paths, drop gettysburg.wav coupling Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- docs/my-website/docs/providers/scaleway.md | 41 +++ .../get_supported_openai_params.py | 9 + .../audio_transcription/transformation.py | 158 ++++++++++++ .../provider_endpoints_support_backup.json | 2 +- litellm/types/utils.py | 1 + litellm/utils.py | 6 + provider_endpoints_support.json | 2 +- ...eway_audio_transcription_transformation.py | 240 ++++++++++++++++++ 8 files changed, 457 insertions(+), 2 deletions(-) create mode 100644 litellm/llms/scaleway/audio_transcription/transformation.py create mode 100644 tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py diff --git a/docs/my-website/docs/providers/scaleway.md b/docs/my-website/docs/providers/scaleway.md index ea57c24db30..8d83a37a3b1 100644 --- a/docs/my-website/docs/providers/scaleway.md +++ b/docs/my-website/docs/providers/scaleway.md @@ -60,3 +60,44 @@ curl http://localhost:4000/chat/completions \ ## Supported features Scaleway provider supports all features in [Generative APIs reference documentation ↗](https://www.scaleway.com/en/developers/api/generative-apis/), such as streaming, structured outputs and tool calling. + +## Audio transcription + +Scaleway's `/audio/transcriptions` endpoint is OpenAI-compatible and works with Whisper models. + +### Python SDK + +```python +import os +from litellm import transcription + +os.environ["SCW_SECRET_KEY"] = "your-scaleway-secret-key" + +with open("speech.mp3", "rb") as audio_file: + response = transcription( + model="scaleway/whisper-large-v3", + file=audio_file, + ) +print(response.text) +``` + +### Proxy config + +```yaml +model_list: + - model_name: scaleway-whisper + litellm_params: + model: scaleway/whisper-large-v3 + api_key: "os.environ/SCW_SECRET_KEY" +``` + +### Proxy request + +```bash +curl http://localhost:4000/v1/audio/transcriptions \ + -H "Authorization: Bearer YOUR_LITELLM_MASTER_KEY" \ + -F model="scaleway-whisper" \ + -F file="@speech.mp3" +``` + +Supported optional params: `language`, `prompt`, `response_format`, `temperature`, `timestamp_granularities`. diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index b72d7abeae0..9d8bd7523db 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -296,6 +296,15 @@ def get_supported_openai_params( # noqa: PLR0915 return OVHCloudAudioTranscriptionConfig().get_supported_openai_params( model=model ) + elif custom_llm_provider == "scaleway": + if request_type == "transcription": + from litellm.llms.scaleway.audio_transcription.transformation import ( + ScalewayAudioTranscriptionConfig, + ) + + return ScalewayAudioTranscriptionConfig().get_supported_openai_params( + model=model + ) elif custom_llm_provider == "elevenlabs": if request_type == "transcription": from litellm.llms.elevenlabs.audio_transcription.transformation import ( diff --git a/litellm/llms/scaleway/audio_transcription/transformation.py b/litellm/llms/scaleway/audio_transcription/transformation.py new file mode 100644 index 00000000000..b45f287afb4 --- /dev/null +++ b/litellm/llms/scaleway/audio_transcription/transformation.py @@ -0,0 +1,158 @@ +""" +Support for Scaleway's OpenAI-compatible `/v1/audio/transcriptions` endpoint. + +API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path-audio-create-an-audio-transcription +""" + +from typing import List, Optional, Union + +import httpx + +from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.llms.base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, + BaseAudioTranscriptionConfig, +) +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIAudioTranscriptionOptionalParams, +) +from litellm.types.utils import FileTypes, TranscriptionResponse + + +class ScalewayAudioTranscriptionException(BaseLLMException): + pass + + +class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIAudioTranscriptionOptionalParams]: + return [ + "language", + "prompt", + "response_format", + "temperature", + "timestamp_granularities", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + for k, v in non_default_params.items(): + if k in supported_params: + optional_params[k] = v + return optional_params + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + api_base = ( + "https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/") + ) + return f"{api_base}/audio/transcriptions" + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return ScalewayAudioTranscriptionException( + message=error_message, + status_code=status_code, + headers=headers, + ) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + if api_key is None: + api_key = get_secret_str("SCW_SECRET_KEY") + + if not api_key: + raise ScalewayAudioTranscriptionException( + message=( + "Scaleway API key not found. Pass `api_key=...` or set the " + "SCW_SECRET_KEY environment variable." + ), + status_code=401, + headers={}, + ) + + default_headers = { + "Authorization": f"Bearer {api_key}", + "accept": "application/json", + } + default_headers.update(headers or {}) + return default_headers + + def transform_audio_transcription_request( + self, + model: str, + audio_file: FileTypes, + optional_params: dict, + litellm_params: dict, + ) -> AudioTranscriptionRequestData: + processed_audio = process_audio_file(audio_file) + + form_fields: dict = {"model": model} + for key in self.get_supported_openai_params(model): + value = optional_params.get(key) + if value is not None: + form_fields[key] = value + + files = { + "file": ( + processed_audio.filename, + processed_audio.file_content, + processed_audio.content_type, + ) + } + + return AudioTranscriptionRequestData(data=form_fields, files=files) + + def transform_audio_transcription_response( + self, + raw_response: httpx.Response, + ) -> TranscriptionResponse: + content_type = (raw_response.headers.get("content-type") or "").lower() + if "application/json" not in content_type: + return TranscriptionResponse(text=raw_response.text) + + try: + response_json = raw_response.json() + except Exception: + raise ScalewayAudioTranscriptionException( + message=raw_response.text, + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + text = response_json.get("text") or "" + response = TranscriptionResponse(text=text) + + if "segments" in response_json: + response["segments"] = response_json["segments"] + if "language" in response_json: + response["language"] = response_json["language"] + + response._hidden_params = response_json + return response diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index ed54c707b00..0562b41d2cd 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1950,7 +1950,7 @@ "responses": true, "embeddings": false, "image_generations": false, - "audio_transcriptions": false, + "audio_transcriptions": true, "audio_speech": false, "moderations": false, "batches": false, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d0bc9b78941..4fe4b124da9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3290,6 +3290,7 @@ class LlmProviders(str, Enum): MANUS = "manus" WANDB = "wandb" OVHCLOUD = "ovhcloud" + SCALEWAY = "scaleway" LEMONADE = "lemonade" AMAZON_NOVA = "amazon_nova" A2A_AGENT = "a2a_agent" diff --git a/litellm/utils.py b/litellm/utils.py index 2125875ee1d..c4aee792972 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8472,6 +8472,12 @@ class ProviderConfigManager: ) return OVHCloudAudioTranscriptionConfig() + elif litellm.LlmProviders.SCALEWAY == provider: + from litellm.llms.scaleway.audio_transcription.transformation import ( + ScalewayAudioTranscriptionConfig, + ) + + return ScalewayAudioTranscriptionConfig() elif litellm.LlmProviders.MISTRAL == provider: from litellm.llms.mistral.audio_transcription.transformation import ( MistralAudioTranscriptionConfig, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 2f3302bb574..6f23c87f911 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1968,7 +1968,7 @@ "responses": true, "embeddings": false, "image_generations": false, - "audio_transcriptions": false, + "audio_transcriptions": true, "audio_speech": false, "moderations": false, "batches": false, diff --git a/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py b/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py new file mode 100644 index 00000000000..407e1d19fb3 --- /dev/null +++ b/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py @@ -0,0 +1,240 @@ +import os +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.scaleway.audio_transcription.transformation import ( + ScalewayAudioTranscriptionConfig, + ScalewayAudioTranscriptionException, +) +from litellm.types.utils import TranscriptionResponse + + +# --------------------------------------------------------------------------- +# get_complete_url +# --------------------------------------------------------------------------- + + +def test_scaleway_get_complete_url_default_base(): + """With no api_base supplied, Scaleway's Generative API endpoint is used.""" + url = ScalewayAudioTranscriptionConfig().get_complete_url( + api_base=None, + api_key="fake", + model="whisper-large-v3", + optional_params={}, + litellm_params={}, + ) + assert url == "https://api.scaleway.ai/v1/audio/transcriptions" + + +def test_scaleway_get_complete_url_custom_base_strips_trailing_slash(): + """Caller-supplied api_base is respected; trailing slash is normalized.""" + url = ScalewayAudioTranscriptionConfig().get_complete_url( + api_base="https://custom.example.com/v1/", + api_key="fake", + model="whisper-large-v3", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.example.com/v1/audio/transcriptions" + + +# --------------------------------------------------------------------------- +# validate_environment +# --------------------------------------------------------------------------- + + +def test_scaleway_validate_environment_explicit_api_key(): + headers = ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + api_key="explicit-key", + ) + assert headers["Authorization"] == "Bearer explicit-key" + assert headers["accept"] == "application/json" + + +def test_scaleway_validate_environment_reads_scw_secret_key(monkeypatch): + monkeypatch.setenv("SCW_SECRET_KEY", "env-secret") + headers = ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + ) + assert headers["Authorization"] == "Bearer env-secret" + + +def test_scaleway_validate_environment_explicit_api_key_wins_over_env(monkeypatch): + """Caller-supplied api_key must win over the SCW_SECRET_KEY env var.""" + monkeypatch.setenv("SCW_SECRET_KEY", "env-secret") + headers = ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + api_key="explicit-wins", + ) + assert headers["Authorization"] == "Bearer explicit-wins" + + +# --------------------------------------------------------------------------- +# transform_audio_transcription_request +# --------------------------------------------------------------------------- + + +def _open_test_audio(): + """Shared helper: open the repo's canonical speech fixture.""" + wav_path = os.path.join( + os.path.dirname(__file__), + "../../../..", + "tests", + "llm_translation", + "gettysburg.wav", + ) + return open(wav_path, "rb") + + +def test_scaleway_transform_request_builds_multipart_with_supported_params(): + with _open_test_audio() as audio_file: + result = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_request( + model="whisper-large-v3", + audio_file=audio_file, + optional_params={ + "language": "en", + "temperature": 0.0, + "response_format": "verbose_json", + }, + litellm_params={}, + ) + ) + + assert isinstance(result.data, dict) + assert result.data["model"] == "whisper-large-v3" + assert result.data["language"] == "en" + assert result.data["temperature"] == 0.0 + assert result.data["response_format"] == "verbose_json" + assert result.files is not None + assert "file" in result.files + assert len(result.files["file"]) == 3 # (filename, content, content_type) + + +def test_scaleway_transform_request_drops_unsupported_params(): + """Only params in get_supported_openai_params() should land in the form.""" + with _open_test_audio() as audio_file: + result = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_request( + model="whisper-large-v3", + audio_file=audio_file, + optional_params={ + "language": "en", + "stream": True, # not supported + "diarize": True, # not supported + }, + litellm_params={}, + ) + ) + + assert "stream" not in result.data + assert "diarize" not in result.data + assert result.data["language"] == "en" + + +# --------------------------------------------------------------------------- +# transform_audio_transcription_response +# --------------------------------------------------------------------------- + + +def test_scaleway_transform_response_parses_text(): + mock_response = MagicMock(spec=httpx.Response) + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = {"text": "Four score and seven years ago"} + + response = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + ) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Four score and seven years ago" + + +def test_scaleway_transform_response_preserves_segments_and_language(): + mock_response = MagicMock(spec=httpx.Response) + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "text": "hello world", + "language": "en", + "segments": [ + {"text": "hello", "start": 0.0, "end": 0.5}, + {"text": "world", "start": 0.6, "end": 1.1}, + ], + } + + response = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + ) + + assert response.text == "hello world" + assert response["language"] == "en" + assert len(response["segments"]) == 2 + + +def test_scaleway_transform_response_raises_typed_exception_on_non_json(): + """Malformed upstream body must raise the Scaleway-typed exception so + error handlers downstream can classify it as a Scaleway failure.""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.side_effect = ValueError("not json") + mock_response.headers = {"content-type": "application/json"} + mock_response.text = "upstream 502 bad gateway" + mock_response.status_code = 502 + + with pytest.raises(ScalewayAudioTranscriptionException): + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + + +def test_scaleway_transform_response_returns_plain_text_for_non_json_content_type(): + """When Scaleway responds with text/srt/vtt (response_format="text" etc.), + the content-type is not application/json — return the body as plain text + rather than exploding on .json().""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.headers = {"content-type": "text/plain; charset=utf-8"} + mock_response.text = "Four score and seven years ago" + + response = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + ) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Four score and seven years ago" + + +def test_scaleway_validate_environment_raises_when_no_key(monkeypatch): + """Missing credential should fail fast with a typed exception rather than + silently emitting 'Bearer None'.""" + monkeypatch.delenv("SCW_SECRET_KEY", raising=False) + + with pytest.raises(ScalewayAudioTranscriptionException) as excinfo: + ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert "SCW_SECRET_KEY" in str(excinfo.value) From 386f334feef5808074de4f201d7d511a6f3acabe Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 20 Apr 2026 15:08:00 -0700 Subject: [PATCH 03/11] Prompt Compression - add it to the proxy (#25729) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor: new agentic loop event hook simplifies how to create logic for tool based multi llm calls * fix: compress - make it work on anthropic input as well * fix(compress.py): working prompt compression for claude code ensures claude code messages can run through proxy easily * docs: add agentic loop hook guide * docs: add agentic_loop_hook to sidebar * fix: fix multiple arguments error * fix: fix tool call loop for compression on streaming /v1/messages * fix: fix linting errors * fix: fix ci/cd errors * feat(litellm_pre_call_utils.py): use claude code session for litellm session id allows claude code logs to be stitched together, making it easy to know they were all part of the same conversation * fix: suppress incorrect mypy warning rE: module * revert: drop PR's changes to litellm/proxy/_experimental/out/ Restores the 34 HTML files under _experimental/out/ to their pre-PR paths (X/index.html -> X.html). All renames are R100 (content unchanged); no other files are touched. * fix: address greptile review comments on PR #25729 - Skip ``kwargs["tools"] = []`` injection when compression is a no-op — Anthropic Messages rejects empty tool arrays on requests that did not originally declare tools. - Move agentic-loop safety guards (fingerprint cycle / max depth) out of the per-callback try/except so they propagate instead of being swallowed by the generic exception handler. Extracted _check_agentic_loop_safety. - Gate generic ``x--session-id`` capture behind the LITELLM_CAPTURE_VENDOR_SESSION_HEADERS env var (off by default) to preserve backwards compatibility; explicit x-litellm-* headers are unaffected. - Fix monkeypatch target in pre-call-hook test to patch the actual module-level binding (litellm.integrations.compression_interception.handler.compress). - Add regression tests for empty-tools skip and opt-in session capture. Co-Authored-By: Claude Opus 4.6 * revert: drop LITELLM_CAPTURE_VENDOR_SESSION_HEADERS flag Generic x--session-id header capture is a new feature and only runs *after* the explicit x-litellm-trace-id / x-litellm-session-id checks, so it does not change behavior for any existing caller that was already using the LiteLLM headers — no backwards-incompatibility to gate. Co-Authored-By: Claude Opus 4.6 * refactor(compress): replace input_type with CallTypes call_type Drop the bespoke ``CompressionInputType`` literal and use the existing ``litellm.types.utils.CallTypes`` enum instead. ``litellm.compress()`` now takes ``call_type: Union[CallTypes, str]`` (default ``CallTypes.completion``) — no new concept to learn, and the enum is already the way the rest of the codebase talks about request shapes. Supported values: ``completion`` / ``acompletion`` (OpenAI chat-completions shape) and ``anthropic_messages`` (Anthropic structured content blocks). Updated: compress(), the compression_interception handler, tests, docs, and the two eval scripts. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- .../docs/completion/prompt_compression.md | 25 + .../docs/proxy/agentic_loop_hook.md | 95 +++ docs/my-website/sidebars.js | 1 + litellm/__init__.py | 1 + litellm/compression/compress.py | 412 +++++++-- .../compression_interception/__init__.py | 14 + .../compression_interception/handler.py | 399 +++++++++ litellm/integrations/custom_logger.py | 37 + .../websearch_interception/handler.py | 336 +++++--- .../messages/agentic_streaming_iterator.py | 320 +++++++ litellm/llms/custom_httpx/llm_http_handler.py | 439 ++++++++-- litellm/proxy/_new_secret_config.yaml | 16 +- litellm/proxy/common_utils/callback_utils.py | 14 + litellm/proxy/litellm_pre_call_utils.py | 57 +- litellm/types/compression.py | 10 +- .../integrations/compression_interception.py | 27 + litellm/types/integrations/custom_logger.py | 30 +- scripts/eval_compression.py | 2 + tests/eval_swe_bench.py | 3 +- .../test_compression_interception_handler.py | 364 ++++++++ .../test_websearch_interception_handler.py | 57 +- .../test_agentic_streaming_iterator.py | 792 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 28 + .../proxy/common_utils/test_callback_utils.py | 35 + .../proxy/test_litellm_pre_call_utils.py | 51 ++ tests/test_litellm/test_compression.py | 325 ++++++- 26 files changed, 3588 insertions(+), 302 deletions(-) create mode 100644 docs/my-website/docs/proxy/agentic_loop_hook.md create mode 100644 litellm/integrations/compression_interception/__init__.py create mode 100644 litellm/integrations/compression_interception/handler.py create mode 100644 litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py create mode 100644 litellm/types/integrations/compression_interception.py create mode 100644 tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py diff --git a/docs/my-website/docs/completion/prompt_compression.md b/docs/my-website/docs/completion/prompt_compression.md index 2d999291af6..0d68ea2c101 100644 --- a/docs/my-website/docs/completion/prompt_compression.md +++ b/docs/my-website/docs/completion/prompt_compression.md @@ -8,6 +8,7 @@ The function keeps high-relevance and recent context, replaces low-relevance con ```python import litellm +from litellm.types.utils import CallTypes messages = [ {"role": "system", "content": "You are a coding assistant."}, @@ -19,6 +20,7 @@ messages = [ compressed = litellm.compress( messages=messages, model="gpt-4o", + call_type=CallTypes.completion, compression_trigger=1000, compression_target=500, ) @@ -45,6 +47,7 @@ response = litellm.completion( - `messages` (`List[dict]`, required): input conversation messages - `model` (`str`, required): model name used for token counting +- `call_type` (`CallTypes`, default `CallTypes.completion`): the LiteLLM call type whose message schema these messages follow. Supported values: `CallTypes.completion` / `CallTypes.acompletion` (OpenAI chat-completions shape) and `CallTypes.anthropic_messages` (Anthropic Messages shape) - `compression_trigger` (`int`, default `200000`): compress only if input token count exceeds this - `compression_target` (`Optional[int]`, default `70% of compression_trigger`): desired post-compression token budget - `embedding_model` (`Optional[str]`): if set, combines BM25 + embedding relevance scoring @@ -70,6 +73,28 @@ args = json.loads(tool_call.function.arguments) full_content = compressed["cache"][args["key"]] ``` +## Server-side Callback Loop (`/v1/messages`) + +You can enable callback-based compression interception to make retrieval loops +transparent for Anthropic Messages calls: + +```yaml +litellm_settings: + callbacks: ["compression_interception"] + compression_interception_params: + enabled: true + compression_trigger: 10000 + compression_target: 7000 +``` + +With this enabled, LiteLLM runs the following server-side flow: + +1. Compresses inbound messages before the first provider call. +2. Injects the `litellm_content_retrieve` tool. +3. Detects retrieval `tool_use` blocks in the model response. +4. Resolves retrieval keys from the compression cache. +5. Reruns the model via agentic loop and returns the final answer. + ## Performance Benchmarked on [SWE-bench Lite](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_bm25_27K) (real GitHub issues with ~27k tokens of BM25-retrieved repo context per problem). diff --git a/docs/my-website/docs/proxy/agentic_loop_hook.md b/docs/my-website/docs/proxy/agentic_loop_hook.md new file mode 100644 index 00000000000..054c03228c4 --- /dev/null +++ b/docs/my-website/docs/proxy/agentic_loop_hook.md @@ -0,0 +1,95 @@ +# Agentic Loop Hook + +Build a `CustomLogger` callback that intercepts a model response, fulfills tool calls server-side, and reruns the model — transparently to the caller. + +:::info Supported call types +- `async` only (sync calls do not trigger the hook) +- Non-streaming only (streaming responses cannot be inspected for tool calls) +- Works on both `/v1/messages` and `/v1/chat/completions` +::: + +## Implement the callback + +Override two methods on `CustomLogger`: + +```python +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch + +MY_TOOL = "my_tool" + +class MyToolCallback(CustomLogger): + + async def async_should_run_agentic_loop( + self, response, model, messages, tools, stream, custom_llm_provider, kwargs + ): + # Return (True, context_dict) if there are tool calls to handle + content = getattr(response, "content", None) or [] + calls = [b for b in content if isinstance(b, dict) + and b.get("type") == "tool_use" and b.get("name") == MY_TOOL] + if not calls: + return False, {} + return True, {"tool_calls": calls} + + async def async_build_agentic_loop_plan( + self, tools, model, messages, response, + anthropic_messages_provider_config, + anthropic_messages_optional_request_params, + logging_obj, stream, kwargs, + ): + calls = tools["tool_calls"] + results = [f"result for {c['input']}" for c in calls] # your logic here + + follow_up = messages + [ + {"role": "assistant", "content": [ + {"type": "tool_use", "id": c["id"], "name": c["name"], "input": c["input"]} + for c in calls + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": c["id"], "content": results[i]} + for i, c in enumerate(calls) + ]}, + ] + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=follow_up), + ) +``` + +For `/v1/chat/completions`, override `async_build_chat_completion_agentic_loop_plan` instead — same idea, `optional_params` replaces `anthropic_messages_optional_request_params`. + +## Register it + +```python +import litellm +litellm.callbacks = [MyToolCallback()] +``` + +Or in `config.yaml`: + +```yaml +litellm_settings: + callbacks: ["my_module.MyToolCallback"] +``` + +## `AgenticLoopPlan` fields + +| Field | Effect | +|---|---| +| `run_agentic_loop=True` + `request_patch` | Reruns the model with the patched request | +| `response_override` | Returns this value directly to the caller (no rerun) | +| `terminate=True` | Stops the loop, returns the current response | +| `run_agentic_loop=False` (default) | Skips; next callback is checked | + +`AgenticLoopRequestPatch` accepts: `model`, `messages`, `tools`, `max_tokens`, `optional_params`, `kwargs`. + +## Loop safety + +- Default max reruns: `3` — override per-request with `kwargs["max_agentic_loops"]` +- Identical tool-call fingerprints abort the loop automatically +- Current depth is in `kwargs["_agentic_loop_depth"]` + +## Examples in this repo + +- `litellm/integrations/compression_interception/handler.py` +- `litellm/integrations/websearch_interception/handler.py` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index c2db54b2237..3d49c142d22 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -536,6 +536,7 @@ const sidebars = { description: "Modify requests, responses, and more", items: [ "proxy/call_hooks", + "proxy/agentic_loop_hook", "proxy/rules", ] }, diff --git a/litellm/__init__.py b/litellm/__init__.py index 3acbb495356..f3bb60c6a09 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -148,6 +148,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "vantage", "posthog", "levo", + "compression_interception", ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index b78b04ec43c..45795c9ca15 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -1,9 +1,9 @@ """ -Main compress() function — orchestrates BM25/embedding scoring, message stubbing, -and retrieval tool injection. +Main compress() function — normalizes input messages, orchestrates BM25/embedding +scoring, message stubbing, and retrieval tool injection. """ -from typing import Any, Dict, List, Optional, Set, Union, cast +from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast from litellm.caching.dual_cache import DualCache from litellm.compression.message_stubbing import ( @@ -15,27 +15,196 @@ from litellm.compression.retrieval_tool import build_retrieval_tool from litellm.compression.scoring.bm25 import bm25_score_messages from litellm.litellm_core_utils.token_counter import token_counter from litellm.types.compression import CompressedResult -from litellm.types.utils import AllMessageValues, Message +from litellm.types.utils import CallTypes + +# CallTypes that produce Anthropic-shaped messages (structured content blocks). +# Everything else is treated as OpenAI chat-completions shape. +_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value}) +# CallTypes that are valid targets for compression. Compression operates on +# message-shaped inputs, so we only accept call types whose payload is a list +# of role/content messages. +_SUPPORTED_CALL_TYPES = frozenset( + { + CallTypes.completion.value, + CallTypes.acompletion.value, + CallTypes.anthropic_messages.value, + } +) + + +def _normalize_call_type(call_type: Union[CallTypes, str]) -> str: + """Return the string value for a ``CallTypes`` enum or a raw string.""" + if isinstance(call_type, CallTypes): + return call_type.value + return call_type + + +def _is_anthropic_call_type(call_type: str) -> bool: + return call_type in _ANTHROPIC_CALL_TYPES + + +def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]: + """ + Build retrieval tool definitions in the target request schema. + + - Chat-completions call types: keep OpenAI function-tool schema. + - Anthropic messages call type: remap to Anthropic's custom tool schema. + """ + if not keys: + return [] + + openai_tools = [build_retrieval_tool(keys)] + if not _is_anthropic_call_type(call_type): + return openai_tools + + # Lazy import to avoid introducing provider transformation imports during + # module import for non-Anthropic call paths. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools) + return cast(List[dict], anthropic_tools) + + +def _content_to_text(content: Any) -> str: + """ + Convert OpenAI/Anthropic message content blocks to plain text. + + Text extraction policy: + - Include text-bearing fields only (`text` blocks + string values). + - For `tool_result`, expand into nested `content` items. + - Ignore non-textual blocks (images/documents/tool metadata/thinking metadata). + + Implemented iteratively (stack-based) to avoid unbounded recursion. + """ + parts: List[str] = [] + stack: List[Any] = [content] + while stack: + item = stack.pop() + if isinstance(item, str): + parts.append(item) + elif isinstance(item, list): + # Push list items in reverse order so they are processed left-to-right. + for element in reversed(item): + stack.append(element) + elif isinstance(item, dict): + item_type = item.get("type") + if item_type == "text": + parts.append(str(item.get("text", ""))) + elif item_type == "tool_result": + stack.append(item.get("content", "")) + return " ".join(parts) + + +def _normalize_messages_for_compression( + messages: List[dict], + call_type: str, +) -> Tuple[List[dict], List[dict]]: + """ + Normalize each original message to a text-surrogate content for scoring. + + Returns: + (normalized_messages, original_messages_copy) + """ + if call_type not in _SUPPORTED_CALL_TYPES: + raise ValueError( + f"Unsupported call_type={call_type!r} for compression. " + f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." + ) + + original_messages: List[Dict[str, Any]] = [dict(m) for m in messages] + + normalized_messages: List[dict] = [] + for msg in original_messages: + normalized_messages.append( + { + **msg, + "content": _content_to_text(msg.get("content", "")), + } + ) + return normalized_messages, original_messages def _extract_last_user_message(messages: List[dict]) -> str: """Return the text content of the last user message.""" for msg in reversed(messages): if msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, str): - return content - if isinstance(content, list): - parts = [] - for part in content: - if isinstance(part, dict) and part.get("type") == "text": - parts.append(part.get("text", "")) - elif isinstance(part, str): - parts.append(part) - return " ".join(parts) + return _content_to_text(msg.get("content", "")) return "" +def _extract_tool_use_ids(content: Any) -> List[str]: + if not isinstance(content, list): + return [] + tool_use_ids: List[str] = [] + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") != "tool_use": + continue + tool_use_id = part.get("id") + if isinstance(tool_use_id, str) and tool_use_id: + tool_use_ids.append(tool_use_id) + return tool_use_ids + + +def _extract_tool_result_ids(content: Any) -> Set[str]: + if not isinstance(content, list): + return set() + tool_result_ids: Set[str] = set() + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") != "tool_result": + continue + tool_use_id = part.get("tool_use_id") + if isinstance(tool_use_id, str) and tool_use_id: + tool_result_ids.add(tool_use_id) + return tool_result_ids + + +def _extract_anthropic_tool_exchange_spans( + messages: List[dict], +) -> Tuple[List[Set[int]], Optional[str]]: + """ + Return atomic 2-message spans for Anthropic tool exchanges. + + Each assistant message containing `tool_use` must be immediately followed by a + user message containing matching `tool_result` blocks for all tool_use ids. + """ + spans: List[Set[int]] = [] + i = 0 + while i < len(messages): + current = messages[i] + if current.get("role") != "assistant": + i += 1 + continue + + tool_use_ids = _extract_tool_use_ids(current.get("content")) + if not tool_use_ids: + i += 1 + continue + + if i + 1 >= len(messages): + return [], "invalid_anthropic_tool_sequence" + + next_msg = messages[i + 1] + if next_msg.get("role") != "user": + return [], "invalid_anthropic_tool_sequence" + + tool_result_ids = _extract_tool_result_ids(next_msg.get("content")) + if not tool_result_ids: + return [], "invalid_anthropic_tool_sequence" + + for tool_use_id in tool_use_ids: + if tool_use_id not in tool_result_ids: + return [], "invalid_anthropic_tool_sequence" + + spans.append({i, i + 1}) + i += 2 + + return spans, None + + def _get_protected_indices(messages: List[dict]) -> List[int]: """ Return indices of messages that must never be compressed: @@ -87,9 +256,98 @@ def _combine_scores( return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)] +def _select_kept_indices_for_budget( + normalized_messages: List[dict], + original_messages: List[dict], + combined_scores: List[float], + compression_target: int, + model: str, + initial_kept_indices: Set[int], + tool_exchange_spans: List[Set[int]], +) -> Tuple[Set[int], Dict[int, dict]]: + kept_indices = set(initial_kept_indices) + current_tokens = 0 + for i in kept_indices: + current_tokens += token_counter( + model=model, + text=cast(str, normalized_messages[i].get("content", "") or ""), + ) + + # Fill token budget from highest-scoring units. + # A unit is either: + # 1) a single message index, or + # 2) an Anthropic tool-exchange span that must be kept/dropped atomically. + truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict + span_id_by_index: Dict[int, int] = {} + for span_id, span in enumerate(tool_exchange_spans): + for idx in span: + span_id_by_index[idx] = span_id + + # Build single-message candidate units (non-span messages). + candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = [] + for idx in range(len(normalized_messages)): + if idx in span_id_by_index or idx in kept_indices: + continue + candidate_units.append((combined_scores[idx], (idx,), True)) + + # Build span candidate units (atomic keep/drop for tool exchanges). + for span in tool_exchange_spans: + span_indices = tuple(sorted(span)) + if any(idx in kept_indices for idx in span_indices): + continue + span_score = max(combined_scores[idx] for idx in span_indices) + candidate_units.append((span_score, span_indices, False)) + + # Sort by descending relevance score. + candidate_units.sort(key=lambda item: item[0], reverse=True) + + for _score, indices, can_truncate in candidate_units: + if any(idx in kept_indices for idx in indices): + continue + msg_tokens = 0 + for idx in indices: + msg_tokens += token_counter( + model=model, + text=cast(str, normalized_messages[idx].get("content", "") or ""), + ) + remaining = compression_target - current_tokens + + if remaining <= 0: + break # budget exhausted + + if current_tokens + msg_tokens <= compression_target: + # Fits entirely + kept_indices.update(indices) + current_tokens += msg_tokens + elif can_truncate and len(indices) == 1 and remaining >= 100: + # Too large to fit whole single message, but we have budget — truncate it. + idx = indices[0] + truncated = truncate_message(original_messages[idx], remaining) + truncated_tokens = token_counter( + model=model, + text=truncated.get("content", "") or "", + ) + truncated_overrides[idx] = truncated + kept_indices.add(idx) + current_tokens += truncated_tokens + + return kept_indices, truncated_overrides + + +def _get_dropped_tool_span_indices( + kept_indices: Set[int], tool_exchange_spans: List[Set[int]] +) -> Set[int]: + dropped_tool_span_indices: Set[int] = set() + for span in tool_exchange_spans: + if not any(idx in kept_indices for idx in span): + dropped_tool_span_indices.update(span) + return dropped_tool_span_indices + + def compress( messages: List[dict], model: str, + call_type: Union[CallTypes, str] = CallTypes.completion, compression_trigger: int = 200_000, compression_target: Optional[int] = None, embedding_model: Optional[str] = None, @@ -108,6 +366,12 @@ def compress( Parameters: messages: The conversation messages to (potentially) compress. model: The LLM model name — used for token counting. + call_type: The LiteLLM call type whose message schema these messages + follow. Supported values: + - ``CallTypes.completion`` / ``CallTypes.acompletion`` — OpenAI + chat-completions shape (default) + - ``CallTypes.anthropic_messages`` — Anthropic Messages shape + (structured content blocks + atomic tool exchanges) compression_trigger: Only compress if input exceeds this token count. compression_target: Target token count after compression. Defaults to ``compression_trigger // 2``. @@ -122,29 +386,37 @@ def compress( A ``CompressedResult`` dict containing compressed messages, token counts, a cache of original content, and the retrieval tool definition. """ + call_type_str = _normalize_call_type(call_type) + normalized_messages, original_messages = _normalize_messages_for_compression( + messages=messages, + call_type=call_type_str, + ) + if compression_target is None: compression_target = compression_trigger * 7 // 10 original_tokens = token_counter( - model=model, messages=cast(List[Union[AllMessageValues, Message]], messages) + model=model, + messages=cast(List[Any], original_messages), ) # Pass through if below trigger if original_tokens <= compression_trigger: return CompressedResult( - messages=messages, + messages=original_messages, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=0.0, cache={}, tools=[], + compression_skipped_reason="below_trigger", ) # Extract query for relevance scoring - query = _extract_last_user_message(messages) + query = _extract_last_user_message(normalized_messages) # Score each message - bm25_scores = bm25_score_messages(query, messages) + bm25_scores = bm25_score_messages(query, normalized_messages) if embedding_model: from litellm.compression.scoring.embedding_scorer import ( @@ -153,7 +425,7 @@ def compress( emb_scores = embedding_score_messages( query, - messages, + normalized_messages, model=embedding_model, cache=compression_cache, embedding_model_params=embedding_model_params, @@ -162,85 +434,69 @@ def compress( else: combined_scores = bm25_scores - # Sort message indices by score descending - ranked_indices = sorted( - range(len(messages)), - key=lambda i: combined_scores[i], - reverse=True, - ) - # Protected messages are never compressed - protected_indices = _get_protected_indices(messages) + protected_indices = _get_protected_indices(normalized_messages) kept_indices: Set[int] = set(protected_indices) - # Count tokens for protected messages - current_tokens = 0 - for i in kept_indices: - current_tokens += token_counter( - model=model, text=messages[i].get("content", "") or "" + tool_exchange_spans: List[Set[int]] = [] + if _is_anthropic_call_type(call_type_str): + tool_exchange_spans, tool_sequence_error = ( + _extract_anthropic_tool_exchange_spans(original_messages) ) - - # Fill token budget from highest-scoring messages. - # For each candidate (ranked by relevance): - # - If it fits entirely → keep it as-is. - # - If it doesn't fit but there's meaningful remaining budget → truncate it - # to fill as much of the budget as possible. - # - Otherwise → stub it (pointer only, content goes to cache). - # Multiple messages may be truncated so we preserve partial content from - # several high-scoring messages rather than fully stubbing all but one. - truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict - - for idx in ranked_indices: - if idx in kept_indices: - continue - msg_content = messages[idx].get("content", "") or "" - msg_tokens = token_counter(model=model, text=msg_content) - remaining = compression_target - current_tokens - - if remaining <= 0: - break # budget exhausted - - if current_tokens + msg_tokens <= compression_target: - # Fits entirely - kept_indices.add(idx) - current_tokens += msg_tokens - elif remaining >= 100: - # Too large to fit whole, but we have budget — truncate it. - truncated = truncate_message(messages[idx], remaining) - truncated_tokens = token_counter( - model=model, - text=truncated.get("content", "") or "", + if tool_sequence_error is not None: + return CompressedResult( + messages=original_messages, + original_tokens=original_tokens, + compressed_tokens=original_tokens, + compression_ratio=0.0, + cache={}, + tools=[], + compression_skipped_reason=tool_sequence_error, ) - truncated_overrides[idx] = truncated - kept_indices.add(idx) - current_tokens += truncated_tokens + + for span in tool_exchange_spans: + # If any message in the span is protected, keep the whole span. + if any(idx in kept_indices for idx in span): + kept_indices.update(span) + + kept_indices, truncated_overrides = _select_kept_indices_for_budget( + normalized_messages=normalized_messages, + original_messages=original_messages, + combined_scores=combined_scores, + compression_target=compression_target, + model=model, + initial_kept_indices=kept_indices, + tool_exchange_spans=tool_exchange_spans, + ) # Build compressed messages and cache compressed_messages: List[dict] = [] cache: Dict[str, str] = {} used_keys: Set[str] = set() + dropped_tool_span_indices = _get_dropped_tool_span_indices( + kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans + ) - for i, msg in enumerate(messages): + for i, msg in enumerate(original_messages): + if i in dropped_tool_span_indices: + continue if i in kept_indices: # Use the truncated version if we made one, otherwise the original compressed_messages.append(truncated_overrides.get(i, msg)) else: - key = extract_key(msg, fallback_index=i, used_keys=used_keys) - content = msg.get("content", "") - if isinstance(content, list): - content = " ".join( - p.get("text", "") if isinstance(p, dict) else str(p) - for p in content - ) + key = extract_key( + normalized_messages[i], fallback_index=i, used_keys=used_keys + ) + content = _content_to_text(msg.get("content", "")) cache[key] = content compressed_messages.append(stub_message(msg, key)) - # Build retrieval tool - tools = [build_retrieval_tool(list(cache.keys()))] if cache else [] + # Build retrieval tool in the target request schema + tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str) compressed_tokens = token_counter( model=model, - messages=cast(List[Union[AllMessageValues, Message]], compressed_messages), + messages=cast(List[Any], compressed_messages), ) return CompressedResult( diff --git a/litellm/integrations/compression_interception/__init__.py b/litellm/integrations/compression_interception/__init__.py new file mode 100644 index 00000000000..14d30af14d8 --- /dev/null +++ b/litellm/integrations/compression_interception/__init__.py @@ -0,0 +1,14 @@ +""" +Compression Interception Module + +Provides server-side prompt compression + retrieval tool fulfillment for +Anthropic Messages agentic loops. +""" + +from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, +) + +__all__ = [ + "CompressionInterceptionLogger", +] diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py new file mode 100644 index 00000000000..c6ae7d9e82b --- /dev/null +++ b/litellm/integrations/compression_interception/handler.py @@ -0,0 +1,399 @@ +""" +Compression Interception Handler + +CustomLogger that compresses inbound Anthropic Messages requests and fulfills +litellm_content_retrieve tool calls server-side via the typed agentic loop plan. +""" + +import time +import uuid +from typing import Any, Dict, List, Optional, Tuple, cast + +from litellm._logging import verbose_logger +from litellm.compression import compress +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.compression_interception import ( + CompressionInterceptionConfig, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import CallTypes + +LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve" +_CACHE_TTL_SECONDS = 15 * 60 + + +class CompressionInterceptionLogger(CustomLogger): + """ + CustomLogger that implements transparent prompt compression + retrieval loops. + + Flow: + 1. Compress inbound /v1/messages requests in pre-call hook. + 2. Inject litellm_content_retrieve tool and persist compressed cache by call_id. + 3. Detect retrieval tool_use blocks in first model response. + 4. Build typed rerun plan with tool_result blocks from the compressed cache. + """ + + def __init__( + self, + enabled: bool = True, + compression_trigger: int = 200_000, + compression_target: Optional[int] = None, + embedding_model: Optional[str] = None, + embedding_model_params: Optional[Dict[str, Any]] = None, + ): + super().__init__() + self.enabled = enabled + self.compression_trigger = compression_trigger + self.compression_target = compression_target + self.embedding_model = embedding_model + self.embedding_model_params = embedding_model_params + self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {} + + @classmethod + def from_config_yaml( + cls, config: CompressionInterceptionConfig + ) -> "CompressionInterceptionLogger": + return cls( + enabled=bool(config.get("enabled", True)), + compression_trigger=int(config.get("compression_trigger", 200_000)), + compression_target=config.get("compression_target"), + embedding_model=config.get("embedding_model"), + embedding_model_params=config.get("embedding_model_params"), + ) + + @staticmethod + def initialize_from_proxy_config( + litellm_settings: Dict[str, Any], + callback_specific_params: Dict[str, Any], + ) -> "CompressionInterceptionLogger": + compression_params: CompressionInterceptionConfig = {} + if "compression_interception_params" in litellm_settings: + compression_params = litellm_settings["compression_interception_params"] + elif "compression_interception" in callback_specific_params: + compression_params = callback_specific_params["compression_interception"] + return CompressionInterceptionLogger.from_config_yaml(compression_params) + + async def async_pre_call_deployment_hook( + self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] + ) -> Optional[dict]: + if not self.enabled: + return None + if call_type is not None and call_type != CallTypes.anthropic_messages: + return None + if int(kwargs.get("_agentic_loop_depth", 0) or 0) > 0: + return None + + messages = kwargs.get("messages") + model = kwargs.get("model") + if not isinstance(messages, list) or not isinstance(model, str): + return None + + if self._has_retrieval_tool(kwargs.get("tools")): + return None + + self._prune_expired_cache() + + compressed = compress( # type: ignore + messages=messages, + model=model, + call_type=CallTypes.anthropic_messages, + compression_trigger=self.compression_trigger, + compression_target=self.compression_target, + embedding_model=self.embedding_model, + embedding_model_params=self.embedding_model_params, + ) + + cache = cast(Dict[str, str], compressed.get("cache", {})) + skip_reason = cast(Optional[str], compressed.get("compression_skipped_reason")) + compressed_tools = cast(List[Dict[str, Any]], compressed.get("tools", [])) + + # Only mutate kwargs when compression actually produced a result. + # If compression was a no-op (below trigger, invalid tool sequence, etc.), + # leave ``messages`` and ``tools`` untouched — injecting an empty + # ``tools: []`` onto a request that originally had no tools breaks + # Anthropic Messages requests. + if cache: + kwargs["messages"] = compressed["messages"] + if compressed_tools: + kwargs["tools"] = self._merge_tools( + existing_tools=cast( + Optional[List[Dict[str, Any]]], kwargs.get("tools") + ), + compressed_tools=compressed_tools, + ) + call_id = cast(Optional[str], kwargs.get("litellm_call_id")) + if not call_id: + call_id = str(uuid.uuid4()) + kwargs["litellm_call_id"] = call_id + self._compression_cache_by_call_id[call_id] = (cache, time.time()) + verbose_logger.debug( + "CompressionInterception: compressed request [call_id=%s original=%d compressed=%d cached_keys=%d]", + call_id, + compressed.get("original_tokens"), + compressed.get("compressed_tokens"), + len(cache), + ) + elif skip_reason is not None: + verbose_logger.debug( + "CompressionInterception: compression skipped [reason=%s original=%d compressed=%d]", + skip_reason, + compressed.get("original_tokens"), + compressed.get("compressed_tokens"), + ) + + return kwargs + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict], + tools: Optional[List[Dict]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict, + ) -> Tuple[bool, Dict]: + if not self.enabled: + return False, {} + if not self._has_retrieval_tool(tools): + return False, {} + + tool_calls, thinking_blocks = self._extract_retrieval_tool_calls( + response=response + ) + if not tool_calls: + return False, {} + + return True, { + "tool_calls": tool_calls, + "thinking_blocks": thinking_blocks, + "tool_type": "compression_retrieval", + } + + async def async_build_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + self._prune_expired_cache() + tool_calls = cast(List[Dict[str, Any]], tools.get("tool_calls", [])) + thinking_blocks = cast(List[Dict[str, Any]], tools.get("thinking_blocks", [])) + + call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) + cache = self._get_cache(call_id=call_id) + retrieval_results = [ + self._resolve_retrieval_content(tc, cache) for tc in tool_calls + ] + + assistant_message = { + "role": "assistant", + "content": thinking_blocks + + [ + { + "type": "tool_use", + "id": tc.get("id"), + "name": tc.get("name", LITELLM_CONTENT_RETRIEVE_TOOL_NAME), + "input": tc.get("input", {}), + } + for tc in tool_calls + ], + } + user_message = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_calls[i].get("id"), + "content": retrieval_results[i], + } + for i in range(len(tool_calls)) + ], + } + follow_up_messages = messages + [assistant_message, user_message] + + max_tokens = cast( + Optional[int], + anthropic_messages_optional_request_params.get("max_tokens") + or kwargs.get("max_tokens"), + ) + optional_params_without_max_tokens = { + k: v + for k, v in anthropic_messages_optional_request_params.items() + if k != "max_tokens" + } + + full_model_name = model + if logging_obj is not None: + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) + full_model_name = cast(str, agentic_params.get("model", model)) + + request_patch = AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + max_tokens=max_tokens, + optional_params=optional_params_without_max_tokens, + kwargs=self._prepare_followup_kwargs(kwargs=kwargs), + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={"tool_type": "compression_retrieval", "call_id": call_id or ""}, + ) + + def _prune_expired_cache(self) -> None: + now = time.time() + self._compression_cache_by_call_id = { + call_id: (cache, created_at) + for call_id, ( + cache, + created_at, + ) in self._compression_cache_by_call_id.items() + if now - created_at <= _CACHE_TTL_SECONDS + } + + def _get_cache(self, call_id: Optional[str]) -> Dict[str, str]: + if not call_id: + return {} + cache_entry = self._compression_cache_by_call_id.get(call_id) + if cache_entry is None: + return {} + return cache_entry[0] + + def _resolve_call_id( + self, logging_obj: Any, kwargs: Dict[str, Any] + ) -> Optional[str]: + if logging_obj is not None: + logging_call_id = getattr(logging_obj, "litellm_call_id", None) + if isinstance(logging_call_id, str) and logging_call_id: + return logging_call_id + kwargs_call_id = kwargs.get("litellm_call_id") + return cast( + Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None + ) + + def _resolve_retrieval_content( + self, tool_call: Dict[str, Any], cache: Dict[str, str] + ) -> str: + raw_input = tool_call.get("input", {}) + key = "" + if isinstance(raw_input, dict): + key = str(raw_input.get("key", "") or "") + if not key: + return "No retrieval key provided." + if key in cache: + return cache[key] + return f"[compressed content key '{key}' not found]" + + def _extract_retrieval_tool_calls( + self, response: Any + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + if isinstance(response, dict): + content = response.get("content", []) + else: + content = getattr(response, "content", []) or [] + + if not isinstance(content, list): + return [], [] + + tool_calls: List[Dict[str, Any]] = [] + thinking_blocks: List[Dict[str, Any]] = [] + + for block in content: + if isinstance(block, dict): + block_type = block.get("type") + block_name = block.get("name") + if block_type in ("thinking", "redacted_thinking"): + thinking_blocks.append(block) + if ( + block_type == "tool_use" + and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ): + tool_calls.append( + { + "id": block.get("id"), + "type": "tool_use", + "name": block_name, + "input": block.get("input", {}), + } + ) + else: + block_type = getattr(block, "type", None) + block_name = getattr(block, "name", None) + if block_type == "thinking": + thinking_blocks.append( + { + "type": "thinking", + "thinking": getattr(block, "thinking", ""), + "signature": getattr(block, "signature", ""), + } + ) + elif block_type == "redacted_thinking": + thinking_blocks.append( + { + "type": "redacted_thinking", + "data": getattr(block, "data", ""), + } + ) + if ( + block_type == "tool_use" + and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ): + tool_calls.append( + { + "id": getattr(block, "id", None), + "type": "tool_use", + "name": block_name, + "input": getattr(block, "input", {}) or {}, + } + ) + + return tool_calls, thinking_blocks + + def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: + internal_keys = {"litellm_logging_obj"} + return { + k: v + for k, v in kwargs.items() + if not k.startswith("_compression_interception") and k not in internal_keys + } + + def _has_retrieval_tool(self, tools: Any) -> bool: + if not isinstance(tools, list): + return False + for tool in tools: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if tool.get("type") == "function" and isinstance(function, dict): + if function.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME: + return True + if ( + tool.get("type") == "custom" + and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ): + return True + return False + + def _merge_tools( + self, + existing_tools: Optional[List[Dict[str, Any]]], + compressed_tools: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + merged = list(existing_tools or []) + if self._has_retrieval_tool(merged): + return merged + merged.extend(compressed_tools) + return merged diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 45c8e2f6262..36486747c39 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER from litellm.types.integrations.argilla import ArgillaItem from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest from litellm.types.prompts.init_prompts import PromptSpec +from litellm.types.integrations.custom_logger import AgenticLoopPlan from litellm.types.utils import ( AdapterCompletionStreamWrapper, CallTypes, @@ -676,6 +677,26 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ pass + async def async_build_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: "LiteLLMLoggingObj", + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + """ + Build a typed rerun plan for Anthropic Messages agentic loops. + + Override this method to separate callback decision/tool execution from + follow-up request execution (handled by BaseLLMHTTPHandler). + """ + return AgenticLoopPlan(run_agentic_loop=False) + async def async_should_run_chat_completion_agentic_loop( self, response: Any, @@ -707,6 +728,22 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ pass + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + optional_params: Dict, + logging_obj: "LiteLLMLoggingObj", + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + """ + Build a typed rerun plan for chat-completions agentic loops. + """ + return AgenticLoopPlan(run_agentic_loop=False) + # Useful helpers for custom logger classes def truncate_standard_logging_payload_content( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 30fd55a3e9d..7b4aa7a3f10 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -28,6 +28,10 @@ from litellm.integrations.websearch_interception.transformation import ( from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -573,6 +577,35 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs=kwargs, ) + async def async_build_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + tool_calls = tools["tool_calls"] + thinking_blocks = tools.get("thinking_blocks", []) + request_patch = await self._build_anthropic_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + thinking_blocks=thinking_blocks, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs, + ) + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={"tool_type": "websearch", "response_format": "anthropic"}, + ) + async def async_run_chat_completion_agentic_loop( self, tools: Dict, @@ -608,6 +641,33 @@ class WebSearchInterceptionLogger(CustomLogger): response_format=response_format, ) + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + optional_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + tool_calls = tools["tool_calls"] + response_format = tools.get("response_format", "openai") + request_patch = await self._build_chat_completion_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + optional_params=optional_params, + kwargs=kwargs, + response_format=response_format, + ) + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={"tool_type": "websearch", "response_format": response_format}, + ) + @staticmethod def _resolve_max_tokens( optional_params: Dict, @@ -672,7 +732,48 @@ class WebSearchInterceptionLogger(CustomLogger): stream: bool, kwargs: Dict, ) -> Any: - """Execute litellm.search() and make follow-up request""" + """Legacy path: execute search + build patch + run follow-up call.""" + request_patch = await self._build_anthropic_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + thinking_blocks=thinking_blocks, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs, + ) + if request_patch.messages is None: + raise ValueError("WebSearchInterception: missing follow-up messages") + + optional_params = dict(anthropic_messages_optional_request_params) + optional_params.update(request_patch.optional_params) + max_tokens = request_patch.max_tokens + if max_tokens is None: + max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None)) + else: + optional_params.pop("max_tokens", None) + if max_tokens is None: + max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + + return await anthropic_messages.acreate( + max_tokens=max_tokens, + messages=request_patch.messages, + model=request_patch.model or model, + **optional_params, + **request_patch.kwargs, + ) + + async def _build_anthropic_request_patch( + self, + model: str, + messages: List[Dict], + tool_calls: List[Dict], + thinking_blocks: List[Dict], + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + kwargs: Dict, + ) -> AgenticLoopRequestPatch: + """Execute litellm.search() and build follow-up request patch.""" # Extract search queries from tool_use blocks search_tasks = [] @@ -721,20 +822,8 @@ class WebSearchInterceptionLogger(CustomLogger): thinking_blocks=thinking_blocks, ) - # Make follow-up request with search results - # Type cast: user_message is a Dict for Anthropic format (default response_format) follow_up_messages = messages + [assistant_message, cast(Dict, user_message)] - verbose_logger.debug( - "WebSearchInterception: Making follow-up request with search results" - ) - verbose_logger.debug( - f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}" - ) - verbose_logger.debug( - f"WebSearchInterception: Last message (tool_result): {user_message}" - ) - # Correlation context for structured logging _call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get( "litellm_call_id", "unknown" @@ -742,61 +831,39 @@ class WebSearchInterceptionLogger(CustomLogger): full_model_name = model # safe default before try block - # Use anthropic_messages.acreate for follow-up request - try: - max_tokens = self._resolve_max_tokens( - anthropic_messages_optional_request_params, kwargs - ) + max_tokens = self._resolve_max_tokens( + anthropic_messages_optional_request_params, kwargs + ) - verbose_logger.debug( - f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request" - ) + verbose_logger.debug( + f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request" + ) - # Create a copy of optional params without max_tokens (since we pass it explicitly) - optional_params_without_max_tokens = { - k: v - for k, v in anthropic_messages_optional_request_params.items() - if k != "max_tokens" - } + optional_params_without_max_tokens = { + k: v + for k, v in anthropic_messages_optional_request_params.items() + if k != "max_tokens" + } + kwargs_for_followup = self._prepare_followup_kwargs(kwargs) - kwargs_for_followup = self._prepare_followup_kwargs(kwargs) - - # Get model from logging_obj.model_call_details["agentic_loop_params"] - # This preserves the full model name with provider prefix (e.g., "bedrock/invoke/...") - if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get( - "agentic_loop_params", {} - ) - full_model_name = agentic_params.get("model", model) - verbose_logger.debug( - f"WebSearchInterception: Using model name: {full_model_name}" - ) - - final_response = await anthropic_messages.acreate( - max_tokens=max_tokens, - messages=follow_up_messages, - model=full_model_name, - **optional_params_without_max_tokens, - **kwargs_for_followup, - ) - verbose_logger.debug( - f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}" - ) - verbose_logger.debug( - f"WebSearchInterception: Final response: {final_response}" - ) - return final_response - except Exception as e: - verbose_logger.exception( - "WebSearchInterception: Follow-up request failed " - "[call_id=%s model=%s messages=%d searches=%d]: %s", - _call_id, - full_model_name, - len(follow_up_messages), - len(final_search_results), - str(e), - ) - raise + if logging_obj is not None: + agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) + full_model_name = agentic_params.get("model", model) + verbose_logger.debug( + "WebSearchInterception: Built anthropic request patch " + "[call_id=%s model=%s messages=%d searches=%d]", + _call_id, + full_model_name, + len(follow_up_messages), + len(final_search_results), + ) + return AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + max_tokens=max_tokens, + optional_params=optional_params_without_max_tokens, + kwargs=kwargs_for_followup, + ) async def _execute_search(self, query: str) -> str: """Execute a single web search using router's search tools""" @@ -883,7 +950,36 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs: Dict, response_format: str = "openai", ) -> Any: - """Execute litellm.search() and make follow-up chat completion request""" + """Legacy path: execute search + build patch + run follow-up call.""" + request_patch = await self._build_chat_completion_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + optional_params=optional_params, + kwargs=kwargs, + response_format=response_format, + ) + if request_patch.messages is None: + raise ValueError("WebSearchInterception: missing follow-up messages") + params = dict(optional_params) + params.update(request_patch.optional_params) + return await litellm.acompletion( + model=request_patch.model or model, + messages=request_patch.messages, + **params, + **request_patch.kwargs, + ) + + async def _build_chat_completion_request_patch( # noqa: PLR0915 + self, + model: str, + messages: List[Dict], + tool_calls: List[Dict], + optional_params: Dict, + kwargs: Dict, + response_format: str = "openai", + ) -> AgenticLoopRequestPatch: + """Execute litellm.search() and build chat-completion rerun patch.""" # Extract search queries from tool_calls search_tasks = [] @@ -963,74 +1059,56 @@ class WebSearchInterceptionLogger(CustomLogger): f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}" ) - # Use litellm.acompletion for follow-up request - try: - # Remove internal parameters that shouldn't be passed to follow-up request - internal_params = { - "_websearch_interception", - "acompletion", - "litellm_logging_obj", - "custom_llm_provider", + # Remove internal parameters that shouldn't be passed to follow-up request + internal_params = { + "_websearch_interception", + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") and k not in internal_params + } + + full_model_name = model + if "custom_llm_provider" in kwargs: + custom_llm_provider = kwargs["custom_llm_provider"] + if not model.startswith(custom_llm_provider) and "/" not in model: + full_model_name = f"{custom_llm_provider}/{model}" + + verbose_logger.debug( + "WebSearchInterception: Built chat completion request patch model=%s messages=%d", + full_model_name, + len(follow_up_messages), + ) + + tools_param = optional_params.get("tools") + optional_params_clean = { + k: v + for k, v in optional_params.items() + if k + not in { + "tools", + "extra_body", "model_alias_map", "stream_response", "custom_prompt_dict", } - kwargs_for_followup = { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and k not in internal_params - } + } + if tools_param is not None: + optional_params_clean["tools"] = tools_param - # Get full model name from kwargs - full_model_name = model - if "custom_llm_provider" in kwargs: - custom_llm_provider = kwargs["custom_llm_provider"] - # Reconstruct full model name with provider prefix if needed - if not model.startswith(custom_llm_provider): - # Check if model already has a provider prefix - if "/" not in model: - full_model_name = f"{custom_llm_provider}/{model}" - - verbose_logger.debug( - f"WebSearchInterception: Using model name: {full_model_name}" - ) - - # Prepare tools for follow-up request (same as original) - tools_param = optional_params.get("tools") - - # Remove tools and extra_body from optional_params to avoid issues - # extra_body often contains internal LiteLLM params that shouldn't be forwarded - optional_params_clean = { - k: v - for k, v in optional_params.items() - if k - not in { - "tools", - "extra_body", - "model_alias_map", - "stream_response", - "custom_prompt_dict", - } - } - - final_response = await litellm.acompletion( - model=full_model_name, - messages=follow_up_messages, - tools=tools_param, - **optional_params_clean, - **kwargs_for_followup, - ) - - verbose_logger.debug( - f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}" - ) - return final_response - except Exception as e: - verbose_logger.exception( - f"WebSearchInterception: Follow-up request failed: {str(e)}" - ) - raise + return AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + optional_params=optional_params_clean, + kwargs=kwargs_for_followup, + ) async def _create_empty_search_result(self) -> str: """Create an empty search result for tool calls without queries""" diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py new file mode 100644 index 00000000000..1f14886ca8e --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -0,0 +1,320 @@ +""" +Agentic Streaming Iterator for Anthropic Messages + +Wraps the raw SSE byte stream from the Anthropic pass-through endpoint, +yields every chunk to the caller (preserving real streaming), collects +all bytes, and on stream exhaustion rebuilds the full Anthropic response +to run through agentic completion hooks. If an agentic hook fires, the +follow-up response is chained as Phase 2 of the same iterator. +""" + +import json +from typing import Any, AsyncIterator, Dict, List, Optional, cast + +from litellm._logging import verbose_logger + + +# --------------------------------------------------------------------------- +# SSE parsing helpers (module-level to keep the class lean) +# --------------------------------------------------------------------------- + + +def _parse_sse_events(raw: bytes) -> List[tuple]: + """Return a list of (event_type, parsed_data_dict) from raw SSE bytes.""" + text = raw.decode("utf-8", errors="replace") + lines = text.split("\n") + events: List[tuple] = [] + current_event_type: Optional[str] = None + + for line in lines: + stripped = line.strip() + if stripped.startswith("event:"): + current_event_type = stripped[len("event:") :].strip() + continue + if not stripped.startswith("data:"): + continue + data_str = stripped[len("data:") :].strip() + try: + data = json.loads(data_str) + except (json.JSONDecodeError, ValueError): + continue + event_type = current_event_type or data.get("type", "") + current_event_type = None + events.append((event_type, data)) + return events + + +def _handle_message_start(data: Dict, response: Dict) -> None: + msg = data.get("message", {}) + response["id"] = msg.get("id", response["id"]) + response["model"] = msg.get("model", response["model"]) + response["role"] = msg.get("role", response["role"]) + usage = msg.get("usage", {}) + if usage: + response["usage"]["input_tokens"] = usage.get("input_tokens", 0) + for key in ("cache_creation_input_tokens", "cache_read_input_tokens"): + if key in usage: + response["usage"][key] = usage[key] + + +def _handle_content_block_start(data: Dict, content_blocks: Dict[int, Dict]) -> None: + idx = data.get("index", len(content_blocks)) + block = data.get("content_block", {}) + block_type = block.get("type", "text") + + _BLOCK_TEMPLATES: Dict[str, Dict] = { + "text": {"type": "text", "text": ""}, + "thinking": {"type": "thinking", "thinking": "", "signature": ""}, + "redacted_thinking": { + "type": "redacted_thinking", + "data": block.get("data", ""), + }, + } + if block_type == "tool_use": + content_blocks[idx] = { + "type": "tool_use", + "id": block.get("id", ""), + "name": block.get("name", ""), + "input": {}, + "_partial_json": "", + } + elif block_type in _BLOCK_TEMPLATES: + content_blocks[idx] = dict(_BLOCK_TEMPLATES[block_type]) + else: + content_blocks[idx] = dict(block) + + +def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> None: + idx = data.get("index", 0) + delta = data.get("delta", {}) + delta_type = delta.get("type", "") + block = content_blocks.get(idx) + if block is None: + return + + if delta_type == "text_delta": + block["text"] = block.get("text", "") + delta.get("text", "") + elif delta_type == "input_json_delta": + block["_partial_json"] = block.get("_partial_json", "") + delta.get( + "partial_json", "" + ) + elif delta_type == "thinking_delta": + block["thinking"] = block.get("thinking", "") + delta.get("thinking", "") + elif delta_type == "signature_delta": + block["signature"] = delta.get("signature", block.get("signature", "")) + + +def _handle_content_block_stop(data: Dict, content_blocks: Dict[int, Dict]) -> None: + idx = data.get("index", 0) + block = content_blocks.get(idx) + if block and block.get("type") == "tool_use": + partial = block.pop("_partial_json", "") + if partial: + try: + block["input"] = json.loads(partial) + except (json.JSONDecodeError, ValueError): + block["input"] = {"_raw": partial} + + +def _handle_message_delta(data: Dict, response: Dict) -> None: + delta = data.get("delta", {}) + if "stop_reason" in delta: + response["stop_reason"] = delta["stop_reason"] + if "stop_sequence" in delta: + response["stop_sequence"] = delta["stop_sequence"] + usage = data.get("usage", {}) + if usage.get("output_tokens") is not None: + response["usage"]["output_tokens"] = usage["output_tokens"] + for key in ( + "input_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ): + if key in usage: + response["usage"][key] = usage[key] + + +class AgenticAnthropicStreamingIterator: + """ + Two-phase async iterator that enables agentic hooks on streaming + Anthropic Messages pass-through responses. + + Phase 1: Yield raw SSE bytes from the upstream response while + accumulating them. When the inner iterator is exhausted, + rebuild the full Anthropic response dict and call agentic hooks. + + Phase 2: If an agentic hook fires and returns a follow-up response + (streaming or non-streaming), yield those bytes to the caller. + """ + + def __init__( + self, + completion_stream: AsyncIterator, + http_handler: Any, + model: str, + messages: List[Dict], + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + custom_llm_provider: str, + kwargs: Dict, + ): + self._inner = completion_stream.__aiter__() + self._http_handler = http_handler + self._model = model + self._messages = messages + self._anthropic_messages_provider_config = anthropic_messages_provider_config + self._anthropic_messages_optional_request_params = ( + anthropic_messages_optional_request_params + ) + self._logging_obj = logging_obj + self._custom_llm_provider = custom_llm_provider + self._kwargs = kwargs + + self._collected_bytes: List[bytes] = [] + self._stream_exhausted = False + self._hook_processing_done = False + self._follow_up_iterator: Optional[AsyncIterator] = None + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + # Phase 1: yield from upstream, collect bytes + if not self._stream_exhausted: + try: + chunk = await self._inner.__anext__() + self._collected_bytes.append(chunk) + return chunk + except StopAsyncIteration: + self._stream_exhausted = True + await self._process_agentic_hooks() + # Fall through to Phase 2 + + # Phase 2: yield from follow-up stream if one was created + if self._follow_up_iterator is not None: + chunk = await self._follow_up_iterator.__anext__() + return chunk + + raise StopAsyncIteration + + async def _process_agentic_hooks(self) -> None: + """Rebuild the Anthropic response from collected SSE bytes and call hooks.""" + if self._hook_processing_done: + return + self._hook_processing_done = True + + if not self._collected_bytes: + return + + try: + rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes) + if rebuilt is None: + verbose_logger.debug( + "AgenticStreamingIterator: Could not rebuild response from SSE bytes" + ) + return + + [ + f"{b.get('type')}({b.get('name', '')})" + if b.get("type") == "tool_use" + else b.get("type") + for b in rebuilt.get("content", []) + ] + + result = await self._http_handler._call_agentic_completion_hooks( + response=rebuilt, + model=self._model, + messages=self._messages, + anthropic_messages_provider_config=self._anthropic_messages_provider_config, + anthropic_messages_optional_request_params=self._anthropic_messages_optional_request_params, + logging_obj=self._logging_obj, + stream=True, + custom_llm_provider=self._custom_llm_provider, + kwargs=self._kwargs, + ) + + if result is None: + return + + if hasattr(result, "__aiter__"): + self._follow_up_iterator = result.__aiter__() + elif isinstance(result, dict): + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) + + fake = FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, result) + ) + self._follow_up_iterator = fake.__aiter__() + else: + verbose_logger.warning( + "AgenticStreamingIterator: Unexpected result type from hooks: %s", + type(result).__name__, + ) + except Exception as e: + _call_id = getattr(self._logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "AgenticStreamingIterator: Error in agentic hook processing " + "[call_id=%s model=%s]: %s", + _call_id, + self._model, + str(e), + ) + + @staticmethod + def _rebuild_anthropic_response_from_sse( + raw_bytes: List[bytes], + ) -> Optional[Dict[str, Any]]: + """ + Parse collected SSE bytes into an Anthropic Messages response dict. + + Processes SSE events in order: + - message_start -> envelope (id, model, role, usage) + - content_block_start -> new content block + - content_block_delta -> accumulate text/json/thinking deltas + - content_block_stop -> finalize block + - message_delta -> stop_reason, output usage + - message_stop -> end + """ + events = _parse_sse_events(b"".join(raw_bytes)) + + response: Dict[str, Any] = { + "id": "", + "type": "message", + "role": "assistant", + "model": "", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + content_blocks: Dict[int, Dict[str, Any]] = {} + saw_message_start = False + + for event_type, data in events: + if event_type == "message_start": + saw_message_start = True + _handle_message_start(data, response) + elif event_type == "content_block_start": + _handle_content_block_start(data, content_blocks) + elif event_type == "content_block_delta": + _handle_content_block_delta(data, content_blocks) + elif event_type == "content_block_stop": + _handle_content_block_stop(data, content_blocks) + elif event_type == "message_delta": + _handle_message_delta(data, response) + + if not saw_message_start: + return None + + for idx in sorted(content_blocks.keys()): + block = content_blocks[idx] + block.pop("_partial_json", None) + response["content"].append(block) + + return response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ea0c05e7656..8a7043111de 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -78,6 +78,10 @@ from litellm.types.containers.main import ( DeleteContainerResult, ) from litellm.types.files import TwoStepFileUploadConfig +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -2047,7 +2051,23 @@ class BaseLLMHTTPHandler: request_body=request_body, litellm_logging_obj=logging_obj, ) - initial_response = completion_stream + + from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, + ) + + initial_response = AgenticAnthropicStreamingIterator( + completion_stream=completion_stream, + http_handler=self, + model=model, + messages=messages, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + return initial_response else: initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response( model=model, @@ -2055,7 +2075,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) - # Call agentic completion hooks + # Call agentic completion hooks (non-streaming path only) final_response = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -2063,7 +2083,7 @@ class BaseLLMHTTPHandler: anthropic_messages_provider_config=anthropic_messages_provider_config, anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, logging_obj=logging_obj, - stream=stream or False, + stream=False, custom_llm_provider=custom_llm_provider, kwargs=kwargs, ) @@ -4516,6 +4536,167 @@ class BaseLLMHTTPHandler: return stream, data return stream, data + @staticmethod + def _get_agentic_loop_settings(kwargs: Dict) -> Tuple[int, int, List[str]]: + depth = int(kwargs.get("_agentic_loop_depth", 0) or 0) + max_loops = int(kwargs.get("max_agentic_loops", 3) or 3) + fingerprints = list(kwargs.get("_agentic_loop_fingerprints", []) or []) + return depth, max(max_loops, 1), fingerprints + + @staticmethod + def _check_agentic_loop_safety( + tool_calls: Any, + fingerprints: List[str], + depth: int, + max_loops: int, + model: str, + ) -> str: + """ + Evaluate agentic-loop safety guards (fingerprint cycle / max depth). + + Raises ValueError on abort. Returns the current fingerprint on success. + + These checks must not be swallowed by the per-callback ``except Exception`` + block that wraps callback dispatch — they are bounded-loop / cycle-break + safety rails and must abort the agentic dispatch when they trip. + """ + fingerprint = BaseLLMHTTPHandler._fingerprint_agentic_tools(tool_calls) + if fingerprint in fingerprints: + raise ValueError( + "Agentic loop detected repeated tool-call fingerprint; aborting rerun" + ) + if depth >= max_loops: + raise ValueError( + f"Exceeded max_agentic_loops={max_loops} for model={model}" + ) + return fingerprint + + @staticmethod + def _fingerprint_agentic_tools(tools: Dict) -> str: + try: + return json.dumps(tools, sort_keys=True, default=str) + except Exception: + return str(tools) + + async def _execute_anthropic_agentic_plan( + self, + plan: AgenticLoopPlan, + model: str, + messages: List[Dict], + anthropic_messages_optional_request_params: Dict, + logging_obj: "LiteLLMLoggingObj", + kwargs: Dict, + depth: int, + max_loops: int, + fingerprints: List[str], + fingerprint: str, + stream: bool = False, + ) -> Any: + from litellm.anthropic_interface import messages as anthropic_messages + + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = model + if logging_obj is not None: + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) + full_model_name = cast(str, agentic_params.get("model", model)) + + optional_params = dict(anthropic_messages_optional_request_params) + optional_params.update(patch.optional_params) + if patch.tools is not None: + optional_params["tools"] = patch.tools + + max_tokens = patch.max_tokens + if max_tokens is None: + max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None)) + else: + optional_params.pop("max_tokens", None) + if max_tokens is None: + max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + + internal_keys = {"litellm_logging_obj"} + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k not in internal_keys + and k not in optional_params + } + kwargs_for_followup.update(patch.kwargs) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + + return await anthropic_messages.acreate( + **{ + "max_tokens": max_tokens, + "messages": patch.messages, + "model": patch.model or full_model_name, + "stream": stream, + **optional_params, + **kwargs_for_followup, + } + ) + + async def _execute_chat_completion_agentic_plan( + self, + plan: AgenticLoopPlan, + model: str, + messages: List[Dict], + optional_params: Dict, + kwargs: Dict, + custom_llm_provider: str, + depth: int, + max_loops: int, + fingerprints: List[str], + fingerprint: str, + ) -> Any: + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = patch.model or model + if "/" not in full_model_name: + full_model_name = f"{custom_llm_provider}/{full_model_name}" + + optional_params_for_followup = dict(optional_params) + optional_params_for_followup.update(patch.optional_params) + if patch.tools is not None: + optional_params_for_followup["tools"] = patch.tools + + internal_params = { + "_websearch_interception", + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k not in internal_params + } + kwargs_for_followup.update(patch.kwargs) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + + return await litellm.acompletion( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) + async def _call_agentic_completion_hooks( self, response: Any, @@ -4541,45 +4722,111 @@ class BaseLLMHTTPHandler: callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) tools = anthropic_messages_optional_request_params.get("tools", []) + depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs) for callback in callbacks: + if not isinstance(callback, CustomLogger): + continue + + should_run: bool = False + tool_calls: Any = None try: - if isinstance(callback, CustomLogger): - # First: Check if agentic loop should run - ( - should_run, - tool_calls, - ) = await callback.async_should_run_agentic_loop( - response=response, + # First: Check if agentic loop should run. Wrap in try/except + # to shield from buggy user callbacks — a callback crash should + # not abort the whole request. + ( + should_run, + tool_calls, + ) = await callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_should_run_agentic_loop [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + continue + + if not should_run: + continue + + # Safety guards must run OUTSIDE the callback try/except — they are + # bounded-loop / cycle-break rails that must propagate to the caller. + fingerprint = self._check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + kwargs_with_provider = kwargs.copy() if kwargs else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + build_plan_overridden = ( + callback.__class__.async_build_agentic_loop_plan + is not CustomLogger.async_build_agentic_loop_plan + ) + if not build_plan_overridden: + return await callback.async_run_agentic_loop( + tools=tool_calls, model=model, messages=messages, - tools=tools, + response=response, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs=kwargs_with_provider, ) - if should_run: - # Second: Execute agentic loop - # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name - kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = ( - custom_llm_provider - ) - agentic_response = await callback.async_run_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - # First hook that runs agentic loop wins - return agentic_response + plan = await callback.async_build_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + verbose_logger.debug( + "Agentic loop terminated by callback=%s reason=%s", + callback.__class__.__name__, + plan.stop_reason, + ) + return response + if not plan.run_agentic_loop: + continue + + return await self._execute_anthropic_agentic_plan( + plan=plan, + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs_with_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + stream=stream, + ) except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( @@ -4653,52 +4900,104 @@ class BaseLLMHTTPHandler: callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) tools = optional_params.get("tools", []) + depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs) for callback in callbacks: - try: - if isinstance(callback, CustomLogger): - # Check if callback has the chat completion agentic loop method - if not hasattr( - callback, "async_should_run_chat_completion_agentic_loop" - ): - continue + if not isinstance(callback, CustomLogger): + continue + if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): + continue - # First: Check if agentic loop should run - ( - should_run, - tool_calls, - ) = await callback.async_should_run_chat_completion_agentic_loop( - response=response, + should_run: bool = False + tool_calls: Any = None + try: + ( + should_run, + tool_calls, + ) = await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_should_run_chat_completion_agentic_loop: %s", + str(e), + ) + continue + + if not should_run: + continue + + # Safety guards must run OUTSIDE the callback try/except — they are + # bounded-loop / cycle-break rails that must propagate to the caller. + fingerprint = self._check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + kwargs_with_provider = kwargs.copy() if kwargs else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + build_plan_overridden = ( + callback.__class__.async_build_chat_completion_agentic_loop_plan + is not CustomLogger.async_build_chat_completion_agentic_loop_plan + ) + if not build_plan_overridden: + return await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, model=model, messages=messages, - tools=tools, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs=kwargs_with_provider, ) - if should_run: - # Second: Execute agentic loop - # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name - kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = ( - custom_llm_provider - ) - agentic_response = ( - await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - ) - # First hook that runs agentic loop wins - return agentic_response + plan = await callback.async_build_chat_completion_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + verbose_logger.debug( + "Agentic chat loop terminated by callback=%s reason=%s", + callback.__class__.__name__, + plan.stop_reason, + ) + return response + if not plan.run_agentic_loop: + continue + + return await self._execute_chat_completion_agentic_plan( + plan=plan, + model=model, + messages=messages, + optional_params=optional_params, + kwargs=kwargs_with_provider, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) except Exception as e: verbose_logger.exception( f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {str(e)}" diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 604e7d5f418..36c90c28559 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -22,11 +22,21 @@ model_list: output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001) # Anthropic model for /v1/messages test — 100x custom pricing - - model_name: "claude-sonnet-4-20250514" + - model_name: "claude-sonnet-4-6" litellm_params: - model: anthropic/claude-sonnet-4-20250514 + model: anthropic/claude-sonnet-4-6 api_key: os.environ/ANTHROPIC_API_KEY model_info: id: claude-sonnet-4-custom-pricing input_cost_per_token: 0.0003 # 100x standard ($0.000003) - output_cost_per_token: 0.0015 # 100x standard ($0.000015) \ No newline at end of file + output_cost_per_token: 0.0015 # 100x standard ($0.000015) + +litellm_settings: + callbacks: ["compression_interception"] + compression_interception_params: + enabled: true + compression_trigger: 100000 +# # optional: +# # embedding_model: "text-embedding-3-small" +# # embedding_model_params: +# # dimensions: 512 \ No newline at end of file diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a206be87a11..e31c76dcac1 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -37,6 +37,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 if isinstance(value, list): imported_list: List[Any] = [] for callback in value: # ["presidio", ] + if isinstance(callback, str) and callback == "compression_interception": + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + compression_interception_obj = ( + CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) + ) + imported_list.append(compression_interception_obj) + continue + # check if callback is a custom logger compatible callback if isinstance(callback, str): callback = LoggingCallbackManager._add_custom_callback_generic_api_str( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 7467bbae232..5804e3f8d9f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1,5 +1,6 @@ import asyncio import copy +import re import time from collections import OrderedDict from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union @@ -28,6 +29,14 @@ _SPECIAL_HEADERS_CACHE = frozenset( v.value.lower() for v in SpecialHeaders._member_map_.values() ) +# Matches any header of the form x--session-id (case-insensitive). +# Excludes the two explicit litellm headers which are handled with higher priority. +_GENERIC_SESSION_ID_HEADER_RE = re.compile(r"^x-.+-session-id$", re.IGNORECASE) +_EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session-id"}) +# Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores +# (covers UUIDs and most common session-id formats). +_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") + def _sanitize_for_log(value: Any) -> str: """ @@ -115,13 +124,43 @@ def _get_metadata_variable_name(request: Request) -> str: return "metadata" +def _extract_generic_session_id_from_headers( + normalized: Dict[str, str], +) -> Optional[str]: + """ + Scan a normalised (lower-cased keys) header dict for any header that looks + like ``x--session-id`` and whose value is a plausible session/trace + identifier (alphanumeric + hyphens/underscores, at least 8 chars). + + The two explicit LiteLLM headers (``x-litellm-trace-id`` / + ``x-litellm-session-id``) are excluded here because they are handled with + higher priority by the caller. + + Example: ``x-claude-code-session-id: e96634a3-fa28-4083-b354-55542e2dca01`` + """ + for key, value in normalized.items(): + if ( + key not in _EXPLICIT_SESSION_HEADERS + and _GENERIC_SESSION_ID_HEADER_RE.match(key) + and isinstance(value, str) + and _SESSION_ID_VALUE_RE.match(value) + ): + return value + return None + + def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]: """ Extract chain id for call chaining from request headers. - x-litellm-trace-id and x-litellm-session-id are interchangeable; when both - are present, x-litellm-trace-id takes precedence. Header keys are matched - case-insensitively so this works with raw header dicts from any transport. + Priority order: + 1. ``x-litellm-trace-id`` (explicit, highest priority) + 2. ``x-litellm-session-id`` (explicit) + 3. Any ``x--session-id`` header whose value looks like a session id + (alphanumeric / UUID, at least 8 chars). E.g. ``x-claude-code-session-id``. + + Header keys are matched case-insensitively so this works with raw header + dicts from any transport. Used by MCP (and other paths that have raw_headers but no Request) to set litellm_trace_id/litellm_session_id for spend logs and logging consistency. @@ -129,8 +168,10 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str if not headers: return None normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)} - return normalized.get("x-litellm-trace-id") or normalized.get( - "x-litellm-session-id" + return ( + normalized.get("x-litellm-trace-id") + or normalized.get("x-litellm-session-id") + or _extract_generic_session_id_from_headers(normalized) ) @@ -649,10 +690,8 @@ class LiteLLMProxyRequestSetup: ######################################################################################### agent_id_from_header = headers.get("x-litellm-agent-id") - # x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining - chain_id = headers.get("x-litellm-trace-id") or headers.get( - "x-litellm-session-id" - ) + # Explicit litellm headers take precedence; fall back to any x-*-session-id header. + chain_id = get_chain_id_from_headers(dict(headers)) if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header diff --git a/litellm/types/compression.py b/litellm/types/compression.py index 01d5a6dd4d6..5dae0c397f0 100644 --- a/litellm/types/compression.py +++ b/litellm/types/compression.py @@ -2,7 +2,14 @@ Type definitions for litellm.compress(). """ -from typing import Dict, List, TypedDict +import sys + +if sys.version_info >= (3, 11): + from typing import Dict, List, NotRequired, TypedDict +else: + from typing import Dict, List, TypedDict + + from typing_extensions import NotRequired class CompressedResult(TypedDict): @@ -12,3 +19,4 @@ class CompressedResult(TypedDict): compression_ratio: float # fraction reduced, e.g. 0.6 means 60% reduction cache: Dict[str, str] # key -> original content (for retrieval tool responses) tools: List[dict] # [litellm_content_retrieve tool definition] + compression_skipped_reason: NotRequired[str] diff --git a/litellm/types/integrations/compression_interception.py b/litellm/types/integrations/compression_interception.py new file mode 100644 index 00000000000..fe52d2ad0d5 --- /dev/null +++ b/litellm/types/integrations/compression_interception.py @@ -0,0 +1,27 @@ +""" +Type definitions for Compression Interception integration. +""" + +from typing import Any, Dict, Optional, TypedDict + + +class CompressionInterceptionConfig(TypedDict, total=False): + """ + Configuration parameters for CompressionInterceptionLogger. + + Used in proxy_config.yaml under litellm_settings: + litellm_settings: + compression_interception_params: + enabled: true + compression_trigger: 100000 + compression_target: 70000 + embedding_model: "text-embedding-3-small" + embedding_model_params: + dimensions: 512 + """ + + enabled: bool + compression_trigger: int + compression_target: Optional[int] + embedding_model: Optional[str] + embedding_model_params: Optional[Dict[str, Any]] diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 06989409229..b5726a11ca0 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -1,6 +1,6 @@ -from typing import Optional +from typing import Any, Dict, List, Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field class StandardCustomLoggerInitParams(BaseModel): @@ -9,3 +9,29 @@ class StandardCustomLoggerInitParams(BaseModel): """ turn_off_message_logging: Optional[bool] = False + + +class AgenticLoopRequestPatch(BaseModel): + """ + Patch returned by callbacks to request a follow-up LLM call. + """ + + model: Optional[str] = None + messages: Optional[List[Dict[str, Any]]] = None + tools: Optional[List[Dict[str, Any]]] = None + max_tokens: Optional[int] = None + optional_params: Dict[str, Any] = Field(default_factory=dict) + kwargs: Dict[str, Any] = Field(default_factory=dict) + + +class AgenticLoopPlan(BaseModel): + """ + Typed callback response for agentic-loop reruns. + """ + + run_agentic_loop: bool = False + request_patch: Optional[AgenticLoopRequestPatch] = None + response_override: Optional[Any] = None + terminate: bool = False + stop_reason: Optional[str] = None + metadata: Dict[str, Any] = Field(default_factory=dict) diff --git a/scripts/eval_compression.py b/scripts/eval_compression.py index d7d90dacc2e..a169cc02d74 100644 --- a/scripts/eval_compression.py +++ b/scripts/eval_compression.py @@ -33,6 +33,7 @@ from dataclasses import asdict, dataclass, field from typing import Optional import litellm +from litellm.types.utils import CallTypes # --------------------------------------------------------------------------- # Problem definitions (HumanEval-style) @@ -880,6 +881,7 @@ def eval_problem( result = litellm.compress( messages=messages, model=model, + call_type=CallTypes.completion, compression_trigger=compression_trigger, embedding_model=embedding_model, ) diff --git a/tests/eval_swe_bench.py b/tests/eval_swe_bench.py index 9c986283abd..6ae99f83ca1 100644 --- a/tests/eval_swe_bench.py +++ b/tests/eval_swe_bench.py @@ -40,6 +40,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import litellm # noqa: E402 from litellm.compression import compress as litellm_compress # noqa: E402 +from litellm.types.utils import CallTypes # noqa: E402 # --------------------------------------------------------------------------- # Prompts @@ -445,7 +446,7 @@ def eval_instance( compress_kwargs: dict = { "messages": messages, "model": model, - "input_type": "openai_chat_completions", + "call_type": CallTypes.completion, "compression_trigger": compression_trigger, "embedding_model": embedding_model, } diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py new file mode 100644 index 00000000000..56e5a94cd49 --- /dev/null +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -0,0 +1,364 @@ +""" +Unit tests for Compression Interception Handler. +""" + +from unittest.mock import MagicMock + +import pytest + +from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, +) +from litellm.types.utils import CallTypes + + +def test_initialize_from_proxy_config(): + """Test initialization from proxy config with litellm_settings.""" + litellm_settings = { + "compression_interception_params": { + "enabled": True, + "compression_trigger": 1234, + "compression_target": 789, + } + } + + logger = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params={}, + ) + + assert logger.enabled is True + assert logger.compression_trigger == 1234 + assert logger.compression_target == 789 + + +@pytest.mark.asyncio +async def test_pre_call_hook_compresses_messages_and_injects_tool(monkeypatch): + """Test pre-call hook compresses and stores per-call cache.""" + logger = CompressionInterceptionLogger() + compressed_result = { + "messages": [{"role": "user", "content": "stubbed"}], + "original_tokens": 12000, + "compressed_tokens": 5000, + "compression_ratio": 0.58, + "cache": {"auth.py": "full file content"}, + "tools": [ + { + "type": "function", + "function": { + "name": "litellm_content_retrieve", + "parameters": { + "type": "object", + "properties": {"key": {"type": "string"}}, + }, + }, + } + ], + } + + def _fake_compress(**kwargs): + return compressed_result + + # The handler does ``from litellm.compression import compress`` at module + # scope, so we must patch the binding on the handler module — patching + # ``litellm.compress`` has no effect on the already-bound reference. + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.compress", + _fake_compress, + ) + + kwargs = { + "model": "bedrock/us.anthropic.claude-sonnet-4-5", + "messages": [{"role": "user", "content": "very large context"}], + "tools": [ + { + "type": "function", + "function": {"name": "existing_tool", "parameters": {"type": "object"}}, + } + ], + } + + result = await logger.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.anthropic_messages + ) + + assert result is not None + assert result["messages"] == compressed_result["messages"] + tool_names = [t.get("function", {}).get("name") for t in result["tools"]] + assert "existing_tool" in tool_names + assert "litellm_content_retrieve" in tool_names + assert result["litellm_call_id"] in logger._compression_cache_by_call_id + + +@pytest.mark.asyncio +async def test_pre_call_hook_below_trigger_does_not_inject_empty_tools(monkeypatch): + """ + When compression is a no-op (below trigger / invalid tool sequence), the + hook must NOT replace ``messages`` or inject an empty ``tools: []`` onto + a request that originally had no tools — Anthropic Messages rejects + ``tools: []``. + """ + logger = CompressionInterceptionLogger() + original_messages = [{"role": "user", "content": "short prompt"}] + + def _fake_compress_noop(**kwargs): + return { + "messages": original_messages, + "original_tokens": 42, + "compressed_tokens": 42, + "compression_ratio": 0.0, + "cache": {}, + "tools": [], + "compression_skipped_reason": "below_trigger", + } + + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.compress", + _fake_compress_noop, + ) + + kwargs = { + "model": "bedrock/us.anthropic.claude-sonnet-4-5", + "messages": original_messages, + } + + result = await logger.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.anthropic_messages + ) + + assert result is not None + # Original request had no ``tools`` — skipped compression must leave it that way. + assert "tools" not in result + # Cache must not be populated for a no-op. + assert result.get("litellm_call_id") not in logger._compression_cache_by_call_id + + +@pytest.mark.asyncio +async def test_should_run_agentic_loop_detects_retrieval_tool_use(): + """Test should-run hook returns tool calls for retrieval tool_use blocks.""" + logger = CompressionInterceptionLogger() + response = { + "content": [ + { + "type": "tool_use", + "id": "toolu_123", + "name": "litellm_content_retrieve", + "input": {"key": "auth.py"}, + } + ] + } + + should_run, tools_dict = await logger.async_should_run_agentic_loop( + response=response, + model="bedrock/claude", + messages=[], + tools=[ + { + "type": "function", + "function": { + "name": "litellm_content_retrieve", + "parameters": {"type": "object"}, + }, + } + ], + stream=False, + custom_llm_provider="bedrock", + kwargs={}, + ) + + assert should_run is True + assert len(tools_dict["tool_calls"]) == 1 + assert tools_dict["tool_calls"][0]["input"]["key"] == "auth.py" + + +@pytest.mark.asyncio +async def test_build_agentic_loop_plan_returns_request_patch(): + """Callback should return typed patch with tool_result content.""" + logger = CompressionInterceptionLogger() + call_id = "call_123" + logger._compression_cache_by_call_id[call_id] = ( + {"auth.py": "full auth file"}, + 9999999999.0, + ) + + logging_obj = MagicMock() + logging_obj.litellm_call_id = call_id + logging_obj.model_call_details = { + "agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"} + } + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_abc", + "type": "tool_use", + "name": "litellm_content_retrieve", + "input": {"key": "auth.py"}, + } + ] + }, + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "read auth.py"}], + response=None, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "tools": [{"name": "litellm_content_retrieve"}], + }, + logging_obj=logging_obj, + stream=False, + kwargs={ + "temperature": 0.1, + "_compression_interception_internal": True, + "litellm_logging_obj": object(), + }, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + assert plan.request_patch.model == "bedrock/invoke/claude-3-5-sonnet" + assert plan.request_patch.max_tokens == 1024 + assert plan.request_patch.messages is not None + assert len(plan.request_patch.messages) == 3 + tool_result_content = plan.request_patch.messages[-1]["content"][0]["content"] + assert tool_result_content == "full auth file" + assert "_compression_interception_internal" not in plan.request_patch.kwargs + assert "litellm_logging_obj" not in plan.request_patch.kwargs + assert plan.request_patch.kwargs["temperature"] == 0.1 + assert "max_tokens" not in plan.request_patch.optional_params + + +@pytest.mark.asyncio +async def test_should_run_agentic_loop_with_custom_type_tools(): + """Test that async_should_run_agentic_loop returns True when tools contain + litellm_content_retrieve as a custom-typed tool (e.g. Claude Code tool list) + and the model response includes a matching tool_use block.""" + logger = CompressionInterceptionLogger() + + # Exact tools payload produced by Claude Code – litellm_content_retrieve is + # the final entry and uses type="custom" (not type="function"). + tools = [ + { + "name": "Agent", + "description": "Launch a new agent to handle complex, multi-step tasks.", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "description": {"type": "string"}, + "prompt": {"type": "string"}, + }, + "required": ["description", "prompt"], + "additionalProperties": False, + }, + }, + { + "name": "AskUserQuestion", + "description": "Use this tool when you need to ask the user questions.", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "questions": {"type": "array", "items": {"type": "object"}}, + }, + "required": ["questions"], + "additionalProperties": False, + }, + }, + { + "name": "Bash", + "description": "Executes a given bash command and returns its output.", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + "additionalProperties": False, + }, + }, + { + "name": "litellm_content_retrieve", + "description": "Retrieve the full content of a file or message that was compressed to save tokens.", + "input_schema": { + "type": "object", + "properties": { + "key": { + "type": "string", + "description": "The identifier of the content to retrieve", + "enum": [ + "message_0", + "HA_UPTIME_ROUTER_SPEC.md", + "message_159", + "message_160", + ], + } + }, + "required": ["key"], + }, + "type": "custom", + }, + ] + + response = { + "content": [ + { + "type": "tool_use", + "id": "toolu_abc", + "name": "litellm_content_retrieve", + "input": {"key": "message_0"}, + } + ] + } + + should_run, tools_dict = await logger.async_should_run_agentic_loop( + response=response, + model="claude-3-5-sonnet", + messages=[], + tools=tools, + stream=False, + custom_llm_provider="anthropic", + kwargs={}, + ) + + assert should_run is True + assert tools_dict["tool_type"] == "compression_retrieval" + assert len(tools_dict["tool_calls"]) == 1 + assert tools_dict["tool_calls"][0]["input"]["key"] == "message_0" + + +@pytest.mark.asyncio +async def test_build_agentic_loop_plan_missing_key_fallback(): + """Missing cache keys should produce deterministic fallback content.""" + logger = CompressionInterceptionLogger() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "missing_call" + logging_obj.model_call_details = {"agentic_loop_params": {}} + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_missing", + "type": "tool_use", + "name": "litellm_content_retrieve", + "input": {"key": "not_found.py"}, + } + ] + }, + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "read file"}], + response=None, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=logging_obj, + stream=False, + kwargs={}, + ) + + assert plan.request_patch is not None + assert ( + plan.request_patch.messages[-1]["content"][0]["content"] + == "[compressed content key 'not_found.py' not found]" + ) diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index c8617a3c1b1..10951265115 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -4,7 +4,7 @@ Unit tests for WebSearch Interception Handler Tests the WebSearchInterceptionLogger class and helper functions. """ -from unittest.mock import MagicMock, Mock +from unittest.mock import AsyncMock, MagicMock, Mock import pytest @@ -69,6 +69,61 @@ async def test_async_should_run_agentic_loop(): assert tools_dict == {} +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_returns_request_patch(): + """Callback should return a typed patch for base handler reruns.""" + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + logger._execute_search = AsyncMock( # type: ignore + return_value="Title: LiteLLM\nURL: docs\nSnippet: test" + ) + + tools_dict = { + "tool_calls": [ + { + "id": "toolu_123", + "type": "tool_use", + "name": "litellm_web_search", + "input": {"query": "what is litellm"}, + } + ], + "response_format": "anthropic", + } + logging_obj = MagicMock() + logging_obj.model_call_details = { + "agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"} + } + kwargs = { + "temperature": 0.2, + "_websearch_interception_converted_stream": True, + "litellm_logging_obj": object(), + } + + plan = await logger.async_build_agentic_loop_plan( + tools=tools_dict, + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "search LiteLLM"}], + response=None, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "tools": [{"name": "litellm_web_search"}], + }, + logging_obj=logging_obj, + stream=False, + kwargs=kwargs, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + assert plan.request_patch.model == "bedrock/invoke/claude-3-5-sonnet" + assert plan.request_patch.max_tokens == 1024 + assert plan.request_patch.messages is not None + assert len(plan.request_patch.messages) == 3 + assert "_websearch_interception_converted_stream" not in plan.request_patch.kwargs + assert "litellm_logging_obj" not in plan.request_patch.kwargs + assert plan.request_patch.kwargs["temperature"] == 0.2 + + @pytest.mark.asyncio async def test_internal_flags_filtered_from_followup_kwargs(): """Test that internal _websearch_interception flags are filtered from follow-up request kwargs. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py new file mode 100644 index 00000000000..b9bda07336f --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -0,0 +1,792 @@ +""" +Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers. +""" + +import json +import os +import sys +from typing import Any, Dict, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, + _handle_content_block_delta, + _handle_content_block_start, + _handle_content_block_stop, + _handle_message_delta, + _handle_message_start, + _parse_sse_events, +) + + +# --------------------------------------------------------------------------- +# Helpers to build SSE byte payloads +# --------------------------------------------------------------------------- + + +def _sse_event(event_type: str, data: dict) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + +def _build_simple_text_stream() -> List[bytes]: + """Produce SSE bytes for a simple text response (no tool calls).""" + chunks = [] + chunks.append( + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello, world!"}, + }, + ) + ) + chunks.append( + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + ) + chunks.append( + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + ) + chunks.append(_sse_event("message_stop", {"type": "message_stop"})) + return chunks + + +def _build_tool_use_stream() -> List[bytes]: + """Produce SSE bytes for a response with a tool_use block.""" + chunks = [] + chunks.append( + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_tool_456", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 50, "output_tokens": 0}, + }, + }, + ) + ) + # thinking block + chunks.append( + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "thinking", + "thinking": "", + "signature": "", + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "thinking_delta", + "thinking": "I need to retrieve...", + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "sig_abc"}, + }, + ) + ) + chunks.append( + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + ) + # tool_use block + chunks.append( + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_001", + "name": "litellm_content_retrieve", + "input": {}, + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": { + "type": "input_json_delta", + "partial_json": '{"key": "section_', + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '1"}'}, + }, + ) + ) + chunks.append( + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 1}) + ) + chunks.append( + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use"}, + "usage": {"output_tokens": 20}, + }, + ) + ) + chunks.append(_sse_event("message_stop", {"type": "message_stop"})) + return chunks + + +# --------------------------------------------------------------------------- +# Mock async stream +# --------------------------------------------------------------------------- + + +class MockAsyncStream: + """Async iterator that yields a list of byte chunks.""" + + def __init__(self, chunks: List[bytes]): + self._chunks = list(chunks) + self._idx = 0 + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if self._idx >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._idx] + self._idx += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Tests for _parse_sse_events +# --------------------------------------------------------------------------- + + +class TestParseSSEEvents: + def test_should_parse_single_event(self): + raw = _sse_event( + "message_start", {"type": "message_start", "message": {"id": "1"}} + ) + events = _parse_sse_events(raw) + assert len(events) == 1 + assert events[0][0] == "message_start" + assert events[0][1]["message"]["id"] == "1" + + def test_should_parse_multiple_events(self): + raw = b"".join(_build_simple_text_stream()) + events = _parse_sse_events(raw) + event_types = [e[0] for e in events] + assert "message_start" in event_types + assert "content_block_start" in event_types + assert "content_block_delta" in event_types + assert "content_block_stop" in event_types + assert "message_delta" in event_types + assert "message_stop" in event_types + + def test_should_skip_malformed_json(self): + raw = b"event: message_start\ndata: {invalid json}\n\n" + events = _parse_sse_events(raw) + assert len(events) == 0 + + def test_should_handle_empty_bytes(self): + events = _parse_sse_events(b"") + assert events == [] + + +# --------------------------------------------------------------------------- +# Tests for _handle_* helpers +# --------------------------------------------------------------------------- + + +class TestHandleMessageStart: + def test_should_populate_envelope(self): + response: Dict[str, Any] = { + "id": "", + "model": "", + "role": "assistant", + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + data = { + "message": { + "id": "msg_abc", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "usage": { + "input_tokens": 42, + "cache_creation_input_tokens": 100, + }, + } + } + _handle_message_start(data, response) + assert response["id"] == "msg_abc" + assert response["model"] == "claude-sonnet-4-20250514" + assert response["usage"]["input_tokens"] == 42 + assert response["usage"]["cache_creation_input_tokens"] == 100 + + +class TestHandleContentBlockStart: + def test_should_create_text_block(self): + blocks: Dict[int, Dict] = {} + data = {"index": 0, "content_block": {"type": "text", "text": ""}} + _handle_content_block_start(data, blocks) + assert blocks[0] == {"type": "text", "text": ""} + + def test_should_create_tool_use_block(self): + blocks: Dict[int, Dict] = {} + data = { + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_x", + "name": "my_tool", + "input": {}, + }, + } + _handle_content_block_start(data, blocks) + assert blocks[1]["type"] == "tool_use" + assert blocks[1]["name"] == "my_tool" + assert blocks[1]["_partial_json"] == "" + + def test_should_create_thinking_block(self): + blocks: Dict[int, Dict] = {} + data = { + "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + } + _handle_content_block_start(data, blocks) + assert blocks[0]["type"] == "thinking" + + +class TestHandleContentBlockDelta: + def test_should_accumulate_text(self): + blocks = {0: {"type": "text", "text": "Hello"}} + _handle_content_block_delta( + {"index": 0, "delta": {"type": "text_delta", "text": " World"}}, + blocks, + ) + assert blocks[0]["text"] == "Hello World" + + def test_should_accumulate_json(self): + blocks = {0: {"type": "tool_use", "_partial_json": '{"key":'}} + _handle_content_block_delta( + { + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '"val"}'}, + }, + blocks, + ) + assert blocks[0]["_partial_json"] == '{"key":"val"}' + + def test_should_ignore_missing_block(self): + blocks: Dict[int, Dict] = {} + _handle_content_block_delta( + {"index": 99, "delta": {"type": "text_delta", "text": "x"}}, + blocks, + ) + assert 99 not in blocks + + +class TestHandleContentBlockStop: + def test_should_parse_tool_input_json(self): + blocks = { + 0: { + "type": "tool_use", + "input": {}, + "_partial_json": '{"key": "section_1"}', + } + } + _handle_content_block_stop({"index": 0}, blocks) + assert blocks[0]["input"] == {"key": "section_1"} + assert "_partial_json" not in blocks[0] + + def test_should_handle_invalid_json_gracefully(self): + blocks = { + 0: { + "type": "tool_use", + "input": {}, + "_partial_json": "not valid json", + } + } + _handle_content_block_stop({"index": 0}, blocks) + assert blocks[0]["input"] == {"_raw": "not valid json"} + + +class TestHandleMessageDelta: + def test_should_set_stop_reason_and_usage(self): + response: Dict[str, Any] = { + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + _handle_message_delta( + { + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 15}, + }, + response, + ) + assert response["stop_reason"] == "end_turn" + assert response["usage"]["output_tokens"] == 15 + + +# --------------------------------------------------------------------------- +# Tests for _rebuild_anthropic_response_from_sse +# --------------------------------------------------------------------------- + + +class TestRebuildAnthropicResponse: + def test_should_rebuild_simple_text_response(self): + raw_bytes = _build_simple_text_stream() + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["id"] == "msg_123" + assert result["model"] == "claude-sonnet-4-20250514" + assert result["stop_reason"] == "end_turn" + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Hello, world!" + assert result["usage"]["input_tokens"] == 10 + assert result["usage"]["output_tokens"] == 5 + + def test_should_rebuild_tool_use_response(self): + raw_bytes = _build_tool_use_stream() + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["id"] == "msg_tool_456" + assert result["stop_reason"] == "tool_use" + assert len(result["content"]) == 2 + + thinking = result["content"][0] + assert thinking["type"] == "thinking" + assert thinking["thinking"] == "I need to retrieve..." + assert thinking["signature"] == "sig_abc" + + tool = result["content"][1] + assert tool["type"] == "tool_use" + assert tool["id"] == "toolu_001" + assert tool["name"] == "litellm_content_retrieve" + assert tool["input"] == {"key": "section_1"} + + def test_should_return_none_without_message_start(self): + raw_bytes = [ + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text"}, + }, + ) + ] + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is None + + def test_should_handle_empty_bytes(self): + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + [] + ) + assert result is None + + def test_should_handle_multi_event_chunks(self): + """When multiple SSE events arrive in a single bytes chunk.""" + combined = b"".join(_build_simple_text_stream()) + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + [combined] + ) + assert result is not None + assert result["content"][0]["text"] == "Hello, world!" + + def test_should_preserve_cache_usage_fields(self): + raw_bytes = [ + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_cache", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "usage": { + "input_tokens": 100, + "cache_creation_input_tokens": 50, + "cache_read_input_tokens": 30, + }, + }, + }, + ), + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 10}, + }, + ), + _sse_event("message_stop", {"type": "message_stop"}), + ] + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["usage"]["cache_creation_input_tokens"] == 50 + assert result["usage"]["cache_read_input_tokens"] == 30 + + def test_should_handle_redacted_thinking_block(self): + raw_bytes = [ + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_redact", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "usage": {"input_tokens": 5}, + }, + }, + ), + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "redacted_thinking", "data": "abc123"}, + }, + ), + _sse_event( + "content_block_stop", + {"type": "content_block_stop", "index": 0}, + ), + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 1}, + }, + ), + _sse_event("message_stop", {"type": "message_stop"}), + ] + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["content"][0]["type"] == "redacted_thinking" + + +# --------------------------------------------------------------------------- +# Tests for AgenticAnthropicStreamingIterator (Phase 1 / Phase 2) +# --------------------------------------------------------------------------- + + +class TestAgenticStreamingIteratorPhase1: + @pytest.mark.asyncio + async def test_should_yield_all_chunks_when_no_hook_fires(self): + """When hooks return None, the wrapper should yield all original chunks.""" + chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert len(collected) == len(chunks) + for orig, got in zip(chunks, collected): + assert orig == got + + mock_handler._call_agentic_completion_hooks.assert_awaited_once() + + @pytest.mark.asyncio + async def test_should_pass_rebuilt_response_to_hooks(self): + """The rebuilt dict passed to hooks should match the original stream content.""" + chunks = _build_tool_use_stream() + mock_stream = MockAsyncStream(chunks) + + captured_response = {} + + async def mock_hooks(**kwargs): + captured_response.update(kwargs["response"]) + return None + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = mock_hooks + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + async for _ in iterator: + pass + + assert captured_response["id"] == "msg_tool_456" + assert captured_response["stop_reason"] == "tool_use" + assert captured_response["content"][1]["name"] == "litellm_content_retrieve" + + +class TestAgenticStreamingIteratorPhase2: + @pytest.mark.asyncio + async def test_should_chain_follow_up_async_iterator(self): + """When hooks return an async iterator, Phase 2 should yield from it.""" + phase1_chunks = _build_simple_text_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_stream = MockAsyncStream(phase1_chunks) + follow_up = MockAsyncStream(phase2_chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=follow_up) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert len(collected) == len(phase1_chunks) + len(phase2_chunks) + assert collected[-2:] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_convert_dict_response_to_fake_stream(self): + """When hooks return a dict, it should be wrapped in FakeAnthropicMessagesStreamIterator.""" + phase1_chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(phase1_chunks) + + fake_response = { + "id": "msg_followup", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [{"type": "text", "text": "follow-up answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 100, "output_tokens": 20}, + } + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=fake_response + ) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + # Phase 1 chunks + Phase 2 fake-stream chunks + assert len(collected) > len(phase1_chunks) + # The follow-up chunks should contain the text from the dict response + phase2_bytes = b"".join(collected[len(phase1_chunks) :]) + assert b"follow-up answer" in phase2_bytes + + +class TestAgenticStreamingIteratorErrorHandling: + @pytest.mark.asyncio + async def test_should_swallow_hook_errors(self): + """Errors in hook processing should be swallowed; Phase 1 chunks are still yielded.""" + chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + side_effect=RuntimeError("hook exploded") + ) + + mock_logging = MagicMock() + mock_logging.litellm_call_id = "test_call_123" + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=mock_logging, + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + # All Phase 1 chunks should still have been yielded + assert len(collected) == len(chunks) + + @pytest.mark.asyncio + async def test_should_handle_empty_stream(self): + """An empty upstream stream should not crash.""" + mock_stream = MockAsyncStream([]) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert collected == [] + # hooks should not be called since no bytes were collected + mock_handler._call_agentic_completion_hooks.assert_not_awaited() + + @pytest.mark.asyncio + async def test_should_pass_stream_true_to_hooks(self): + """The wrapper should always pass stream=True to hooks.""" + chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + async for _ in iterator: + pass + + call_kwargs = mock_handler._call_agentic_completion_hooks.call_args + assert call_kwargs.kwargs["stream"] is True diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 2b9d2e9e543..6924eb8d3d9 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -74,6 +74,34 @@ def test_prepare_fake_stream_request(): assert result_data["messages"] == [{"role": "user", "content": "Hello"}] +def test_get_agentic_loop_settings_defaults_and_overrides(): + handler = BaseLLMHTTPHandler() + + depth, max_loops, fingerprints = handler._get_agentic_loop_settings(kwargs={}) + assert depth == 0 + assert max_loops == 3 + assert fingerprints == [] + + depth, max_loops, fingerprints = handler._get_agentic_loop_settings( + kwargs={ + "_agentic_loop_depth": 2, + "max_agentic_loops": 7, + "_agentic_loop_fingerprints": ["fp-1", "fp-2"], + } + ) + assert depth == 2 + assert max_loops == 7 + assert fingerprints == ["fp-1", "fp-2"] + + +def test_fingerprint_agentic_tools_is_deterministic(): + handler = BaseLLMHTTPHandler() + tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]} + tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]} + + assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b) + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_extra_headers(): """ diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 27528fbd20b..c6132194c74 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,14 +1,17 @@ import sys import os +from types import SimpleNamespace sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path from litellm.proxy.common_utils.callback_utils import ( + initialize_callbacks_on_proxy, get_remaining_tokens_and_requests_from_request_data, normalize_callback_names, ) +import litellm from unittest.mock import patch from litellm.proxy.common_utils.callback_utils import process_callback @@ -84,3 +87,35 @@ def test_normalize_callback_names_lowercases_strings(): "s3", "custom_callback", ] + + +def test_initialize_callbacks_on_proxy_instantiates_compression_interception( + monkeypatch, +): + dummy_callback = object() + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.CompressionInterceptionLogger.initialize_from_proxy_config", + lambda litellm_settings, callback_specific_params: dummy_callback, + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + initialize_callbacks_on_proxy( + value=["compression_interception"], + premium_user=False, + config_file_path=".", + litellm_settings={"compression_interception_params": {"enabled": True}}, + callback_specific_params={}, + ) + assert dummy_callback in litellm.callbacks + assert "compression_interception" not in litellm.callbacks + finally: + litellm.callbacks = original_callbacks diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 34d3c203377..ac009df67b0 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1791,6 +1791,57 @@ def test_add_litellm_metadata_from_request_headers_both_headers_trace_id_precede assert data["litellm_trace_id"] == "trace-value" +def test_add_litellm_metadata_from_request_headers_generic_session_id_header(): + """A generic x--session-id header is used when no explicit litellm header is set.""" + headers = {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert data["litellm_session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert data["litellm_trace_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + + +def test_add_litellm_metadata_from_request_headers_explicit_header_beats_generic(): + """Explicit x-litellm-trace-id wins over a generic x-*-session-id header.""" + headers = { + "x-litellm-trace-id": "explicit-trace-id-value", + "x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_session_id"] == "explicit-trace-id-value" + assert data["litellm_trace_id"] == "explicit-trace-id-value" + + +def test_get_chain_id_from_headers_generic_vendor_session_id(): + """get_chain_id_from_headers picks up any x--session-id with a valid value.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert ( + get_chain_id_from_headers( + {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} + ) + == "e96634a3-fa28-4083-b354-55542e2dca01" + ) + # Short / non-alphanumeric values should be ignored + assert get_chain_id_from_headers({"x-foo-session-id": "short"}) is None + assert get_chain_id_from_headers({"x-foo-session-id": "has spaces!!"}) is None + # Explicit headers still take precedence + assert ( + get_chain_id_from_headers( + { + "x-litellm-trace-id": "explicit-id-value", + "x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01", + } + ) + == "explicit-id-value" + ) + + def test_get_internal_user_header_from_mapping_returns_expected_header(): mappings = [ {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, diff --git a/tests/test_litellm/test_compression.py b/tests/test_litellm/test_compression.py index 13dda0cbcbc..4fbcd4ed30d 100644 --- a/tests/test_litellm/test_compression.py +++ b/tests/test_litellm/test_compression.py @@ -3,6 +3,7 @@ Unit tests for litellm.compress(). """ import os +import importlib import pytest @@ -12,6 +13,10 @@ from litellm.compression.scoring.embedding_scorer import embedding_score_message from litellm.compression.content_detection import detect_content_type from litellm.compression.message_stubbing import extract_key, stub_message from litellm.compression.retrieval_tool import build_retrieval_tool +from litellm.types.utils import CallTypes + +CALL_TYPE = CallTypes.completion +ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- @@ -149,7 +154,7 @@ def test_retrieval_tool_description_lists_keys(): def test_compress_below_trigger_passthrough(): messages = [{"role": "user", "content": "hello"}] - result = litellm.compress(messages, model="gpt-4o") + result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) assert result["messages"] == messages assert result["cache"] == {} assert result["tools"] == [] @@ -178,6 +183,7 @@ def test_compress_above_trigger(): result = litellm.compress( big_messages, model="gpt-4o", + call_type=CALL_TYPE, compression_trigger=1000, compression_target=500, ) @@ -189,13 +195,62 @@ def test_compress_above_trigger(): assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve" +def test_compress_anthropic_list_content_is_boundary_stable(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System prompt"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "# a.py\n" + "alpha " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/a.png"}, + }, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "# b.py\n" + "beta " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/b.png"}, + }, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "Fix alpha bug in a.py"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] < result["original_tokens"] + assert len(result["messages"]) == len(messages) + assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages] + assert len(result["cache"]) > 0 + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "custom" + assert result["tools"][0]["name"] == "litellm_content_retrieve" + assert "input_schema" in result["tools"][0] + + def test_compress_preserves_system_message(): messages = [ {"role": "system", "content": "System prompt. " * 500}, {"role": "user", "content": "Large file content. " * 5000}, {"role": "user", "content": "Fix the bug"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) assert result["messages"][0]["role"] == "system" assert "System prompt" in result["messages"][0]["content"] @@ -205,7 +260,9 @@ def test_compress_preserves_last_user_message(): {"role": "user", "content": "Big context " * 5000}, {"role": "user", "content": "Fix the bug in auth.py"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) last_user = [m for m in result["messages"] if m["role"] == "user"][-1] assert "Fix the bug in auth.py" in last_user["content"] @@ -216,7 +273,9 @@ def test_compress_preserves_last_assistant_message(): {"role": "assistant", "content": "I'll help with that. " * 2000}, {"role": "user", "content": "Now fix the bug"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] assert len(assistant_msgs) >= 1 # The last assistant message should be preserved (not stubbed) @@ -229,7 +288,9 @@ def test_cache_keys_match_stubs(): {"role": "user", "content": "# auth.py\n" + "code " * 5000}, {"role": "user", "content": "Fix it"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) if result["tools"]: tool_desc = result["tools"][0]["function"]["description"] for key in result["cache"]: @@ -242,11 +303,75 @@ def test_compress_default_target(): {"role": "user", "content": "content " * 5000}, {"role": "user", "content": "query"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=2000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 + ) # Should have compressed — target = 1000 assert result["compressed_tokens"] <= result["original_tokens"] +def test_compress_nested_tool_result_extracts_text_only(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System rules"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "prefix"}, + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "nested text fragment"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/secret-tool.png", + }, + }, + ], + }, + { + "type": "image_url", + "image_url": {"url": "https://example.com/top.png"}, + }, + {"type": "text", "text": " " + ("irrelevant " * 3000)}, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "final query that must remain"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=500, + compression_target=100, + ) + + cached_text = " ".join(result["cache"].values()) + assert "nested text fragment" in cached_text + assert "https://example.com/secret-tool.png" not in cached_text + assert "https://example.com/top.png" not in cached_text + + +def test_compress_default_call_type_is_completion(): + result = litellm.compress( + messages=[ + {"role": "user", "content": "Large context " * 4000}, + {"role": "user", "content": "query"}, + ], + model="gpt-4o", + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] <= result["original_tokens"] + assert isinstance(result["tools"], list) + + def test_compress_forwards_embedding_model_params(monkeypatch): captured = {} @@ -269,6 +394,7 @@ def test_compress_forwards_embedding_model_params(monkeypatch): {"role": "user", "content": "Fix auth"}, ], model="gpt-4o", + call_type=CALL_TYPE, compression_trigger=1000, embedding_model="text-embedding-3-small", embedding_model_params={"api_base": "https://example-embeddings.test"}, @@ -326,6 +452,7 @@ def test_embedding_scorer(): {"role": "user", "content": "Fix auth"}, ], model="gpt-4o", + call_type=CALL_TYPE, compression_trigger=1000, embedding_model="text-embedding-3-small", ) @@ -346,8 +473,9 @@ def test_simple_compression(final_user_message, expected_content): {"role": "user", "content": "Unrelated cooking recipes " * 2000}, {"role": "user", "content": final_user_message}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) - print(result["messages"]) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) if expected_content == "Unrelated cooking recipes ": assert "Unrelated cooking recipes " in result["messages"][1]["content"] assert "Authentication code " not in result["messages"][0]["content"] @@ -356,3 +484,184 @@ def test_simple_compression(final_user_message, expected_content): assert "Unrelated cooking recipes " not in result["messages"][1]["content"] else: raise ValueError(f"Unexpected expected_content: {expected_content}") + + +def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2) + return [0.95, 0.01, 0.02, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_drop", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_drop", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + # idx=1,2 should be dropped atomically (no orphan tool blocks left behind) + assert len(result["messages"]) == 3 + assert result["messages"][0]["role"] == "user" + assert "other_blob" in result["messages"][0]["content"] + assert result["messages"][1]["content"] == "assistant_tail" + assert result["messages"][2]["content"] == "final query" + assert result["cache"] == {} + + +def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer the tool exchange span over idx=0 + return [0.05, 0.01, 0.92, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_keep", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_keep", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert len(result["messages"]) == 5 + assert result["messages"][1]["role"] == "assistant" + assert result["messages"][2]["role"] == "user" + # idx=0 should be compressed instead + assert "litellm_content_retrieve" in result["messages"][0]["content"] + assert len(result["cache"]) == 1 + + +def test_compress_anthropic_malformed_tool_sequence_passes_through(): + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_broken", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert result["messages"] == messages + assert result["cache"] == {} + assert result["tools"] == [] + assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence" From 4f823cedac47473ea0dae58ccbe10d0afbb725d2 Mon Sep 17 00:00:00 2001 From: Michael-RZ-Berri Date: Mon, 20 Apr 2026 15:25:21 -0700 Subject: [PATCH 04/11] Add supported providers to prompt caching doc (#26124) * Add supported providers to prompt caching doc * Move Z.ai / GLM to cache_control marker list * Mark xAI models as supporting prompt caching * Narrow xAI prompt caching flag to models with documented cache pricing * Add prompt caching flag to grok-4, grok-4-0709, grok-4-latest --------- Co-authored-by: Michael Riad Zaky --- .../docs/completion/prompt_caching.md | 1 + .../docs/tutorials/prompt_caching.md | 16 ++++++++ model_prices_and_context_window.json | 39 ++++++++++++++++--- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/docs/my-website/docs/completion/prompt_caching.md b/docs/my-website/docs/completion/prompt_caching.md index 402c7b9f4c7..aaae7e7be76 100644 --- a/docs/my-website/docs/completion/prompt_caching.md +++ b/docs/my-website/docs/completion/prompt_caching.md @@ -10,6 +10,7 @@ Supported Providers: - Vertex AI (`vertex_ai/`, `vertex_ai_beta/`) - Bedrock (`bedrock/`, `bedrock/invoke/`, `bedrock/converse`) ([All models bedrock supports prompt caching on](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html)) - Deepseek API (`deepseek/`) +- xAI (`xai/`) For the supported providers, LiteLLM follows the OpenAI prompt caching usage object format: diff --git a/docs/my-website/docs/tutorials/prompt_caching.md b/docs/my-website/docs/tutorials/prompt_caching.md index ab2aa00d773..581d2ba7c36 100644 --- a/docs/my-website/docs/tutorials/prompt_caching.md +++ b/docs/my-website/docs/tutorials/prompt_caching.md @@ -8,6 +8,22 @@ Reduce costs by up to 90% by using LiteLLM to auto-inject prompt caching checkpo +Supported Providers (`cache_control` marker): +- Anthropic API (`anthropic/`) +- AWS Bedrock - Claude (`bedrock/`) +- Vertex AI - Claude and Gemini (`vertex_ai/`) +- Google AI Studio - Gemini (`gemini/`) +- Azure AI - Claude (`azure_ai/`) +- OpenRouter - Claude, Gemini, MiniMax, GLM, z-ai routes (`openrouter/`) +- Databricks - Claude (`databricks/`) +- DashScope / Qwen (`dashscope/`) +- MiniMax (`minimax/`) +- Z.ai / GLM (`zai/`) + +Provider Managed (automatic, no marker needed): +- OpenAI (`openai/`) +- DeepSeek (`deepseek/`) +- xAI (`xai/`) ## How it works diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 04b68b8f4ec..72806369ea5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33302,6 +33302,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33317,6 +33318,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33332,6 +33334,7 @@ "output_cost_per_token": 2.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33347,6 +33350,7 @@ "output_cost_per_token": 2.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33362,6 +33366,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33378,6 +33383,7 @@ "output_cost_per_token": 5e-07, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33395,6 +33401,7 @@ "output_cost_per_token": 5e-07, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33411,6 +33418,7 @@ "output_cost_per_token": 4e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33427,6 +33435,7 @@ "output_cost_per_token": 4e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33443,6 +33452,7 @@ "output_cost_per_token": 4e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33459,6 +33469,7 @@ "output_cost_per_token": 5e-07, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33474,38 +33485,41 @@ "output_cost_per_token": 1.5e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, "xai/grok-4-fast-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, - "cache_read_input_token_cost": 5e-08, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, "xai/grok-4-fast-non-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, - "cache_read_input_token_cost": 5e-08, "max_tokens": 2000000.0, "mode": "chat", - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, @@ -33521,6 +33535,7 @@ "output_cost_per_token_above_128k_tokens": 3e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, @@ -33536,6 +33551,7 @@ "output_cost_per_token_above_128k_tokens": 3e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, @@ -33553,6 +33569,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -33573,6 +33590,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -33593,6 +33611,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -33613,6 +33632,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -33632,6 +33652,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -33648,6 +33669,7 @@ "output_cost_per_token": 6e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, @@ -33664,6 +33686,7 @@ "output_cost_per_token": 6e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, @@ -33696,6 +33719,7 @@ "output_cost_per_token": 6e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true @@ -33724,6 +33748,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, @@ -33738,6 +33763,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, @@ -33752,6 +33778,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, From 9aee0da7d8c3618b6505a83a646ab388bef77745 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 20 Apr 2026 15:29:43 -0700 Subject: [PATCH 05/11] fix: /health/readiness 503 loop when DB is unreachable (#26134) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: /health/readiness returns 503 when DB is unreachable due to handle_db_exception re-raising handle_db_exception() re-raises the Prisma exception inside _db_health_readiness_check's except block, which propagates out to health_readiness() and gets wrapped in a 503. The health endpoint never reached the reconnect path and the service never recovered. Fix: - Remove handle_db_exception() call from _db_health_readiness_check — that helper is for API request handlers (allow_requests_on_db_unavailable flag), not health checks - Replace raw disconnect()+connect() with attempt_db_reconnect(), which uses the proper lock, cooldown, escalation, and heavy-reconnect (recreate_prisma_client) machinery * test: update health readiness tests for handle_db_exception removal - Remove tests that expected handle_db_exception to re-raise (old buggy behaviour) - Remove tests asserting disconnect()/connect() calls (replaced by attempt_db_reconnect) - Add regression tests covering the 503 loop fix: - transport errors never raise (ClientNotConnectedError, httpx.ConnectError, etc.) - reconnect success path returns 'connected' - reconnect failure path returns 'disconnected' without raising - non-transport errors return 'disconnected', skip reconnect --------- Co-authored-by: yuneng-jiang --- .../health_endpoints/_health_endpoints.py | 6 +- .../health_endpoints/test_health_endpoints.py | 218 +++++------------- 2 files changed, 66 insertions(+), 158 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8fd19548cbb..b4b5de1746e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1121,14 +1121,14 @@ async def _db_health_readiness_check(): return db_health_cache except Exception as e: db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} - PrismaDBExceptionHandler.handle_db_exception(e) if PrismaDBExceptionHandler.is_database_transport_error(e): try: verbose_proxy_logger.warning( "_db_health_readiness_check: health_check failed, attempting reconnect" ) - await prisma_client.disconnect() - await prisma_client.connect() + await prisma_client.attempt_db_reconnect( + reason="health_readiness_check" + ) await prisma_client.health_check() verbose_proxy_logger.info( "_db_health_readiness_check: reconnect succeeded" diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index c275c665114..ba260142351 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -9,6 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +import httpx import pytest from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError @@ -110,220 +111,127 @@ async def test_db_health_prisma_client_none(): @pytest.mark.asyncio @pytest.mark.parametrize( - "prisma_error", + "transport_error", [ - PrismaError(), + httpx.ConnectError("All connection attempts failed"), ClientNotConnectedError(), HTTPClientClosedError(), + PrismaError("Can't reach database server"), ], ) -async def test_db_health_error_flag_off_raises_no_reconnect(prisma_error): +async def test_db_health_transport_error_never_raises(transport_error): """ - When health_check raises and allow_requests_on_db_unavailable is False, - handle_db_exception re-raises immediately. The reconnect path is never - reached, so disconnect/connect are never called. + Regression test for the /health/readiness 503 loop bug. + + handle_db_exception() used to re-raise inside _db_health_readiness_check, + turning any DB outage into a 503 "Service Unhealthy" response that never + recovered. Transport errors (ClientNotConnectedError, httpx.ConnectError, + etc.) must return {"status": "disconnected"} — never raise. """ mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=prisma_error) - mock_prisma.disconnect = AsyncMock() + mock_prisma.health_check = AsyncMock(side_effect=transport_error) + mock_prisma.attempt_db_reconnect = AsyncMock(return_value=False) _health_endpoints_module.db_health_cache = { "status": "connected", "last_updated": datetime.now() - timedelta(seconds=20), } - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": False}, - ), - ): - with pytest.raises(Exception) as exc_info: - await _db_health_readiness_check() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await _db_health_readiness_check() - assert exc_info.value is prisma_error - mock_prisma.disconnect.assert_not_called() - assert _health_endpoints_module.db_health_cache["status"] == "disconnected" + assert result["status"] == "disconnected" + mock_prisma.attempt_db_reconnect.assert_called_once_with( + reason="health_readiness_check" + ) @pytest.mark.asyncio @pytest.mark.parametrize( - "prisma_error", + "transport_error", [ - PrismaError("Can't reach database server"), + httpx.ConnectError("All connection attempts failed"), ClientNotConnectedError(), HTTPClientClosedError(), ], ) -async def test_db_health_error_flag_on_reconnect_succeeds(prisma_error): +async def test_db_health_transport_error_reconnect_succeeds(transport_error): """ - When health_check raises, allow_requests_on_db_unavailable is True, - and the reconnect cycle (disconnect -> connect -> health_check) succeeds, - return 'connected' and update the cache. + When health_check raises a transport error and attempt_db_reconnect + succeeds, the second health_check passes and we return 'connected'. """ mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=[prisma_error, None]) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() + mock_prisma.health_check = AsyncMock(side_effect=[transport_error, None]) + mock_prisma.attempt_db_reconnect = AsyncMock(return_value=True) _health_endpoints_module.db_health_cache = { "status": "connected", "last_updated": datetime.now() - timedelta(seconds=20), } - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): result = await _db_health_readiness_check() assert result["status"] == "connected" - mock_prisma.disconnect.assert_called_once() - mock_prisma.connect.assert_called_once() + mock_prisma.attempt_db_reconnect.assert_called_once_with( + reason="health_readiness_check" + ) assert mock_prisma.health_check.call_count == 2 @pytest.mark.asyncio @pytest.mark.parametrize( - "prisma_error", + "transport_error", [ - PrismaError("Can't reach database server"), + httpx.ConnectError("All connection attempts failed"), ClientNotConnectedError(), HTTPClientClosedError(), ], ) -async def test_db_health_error_flag_on_reconnect_fails(prisma_error): +async def test_db_health_transport_error_reconnect_fails(transport_error): """ - When health_check raises, allow_requests_on_db_unavailable is True, - but the reconnect also fails, return 'disconnected' instead of raising. - This respects the flag's intent: keep serving even without a DB. + When health_check raises a transport error and attempt_db_reconnect also + fails, return 'disconnected' without raising. """ mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=prisma_error) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): - result = await _db_health_readiness_check() - - assert result["status"] == "disconnected" - mock_prisma.disconnect.assert_called_once() - mock_prisma.connect.assert_called_once() - - -@pytest.mark.asyncio -async def test_db_health_non_transport_error_flag_off_raises(): - """ - When health_check raises a non-transport error and - allow_requests_on_db_unavailable is False, handle_db_exception - re-raises before reaching the is_database_transport_error guard. - Cache is still invalidated before the re-raise. - """ - non_transport_error = PrismaError("UniqueViolationError") - mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=non_transport_error) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": False}, - ), - ): - with pytest.raises(PrismaError): - await _db_health_readiness_check() - - assert _health_endpoints_module.db_health_cache["status"] == "disconnected" - mock_prisma.disconnect.assert_not_called() - mock_prisma.connect.assert_not_called() - - -@pytest.mark.asyncio -async def test_db_health_non_transport_error_flag_on_skips_reconnect(): - """ - When health_check raises a non-transport error (e.g. data-layer) and - allow_requests_on_db_unavailable is True, handle_db_exception swallows - the exception, then is_database_transport_error returns False so the - reconnect cycle is skipped. Returns 'disconnected' without calling - disconnect/connect. - """ - non_transport_error = PrismaError("UniqueViolationError") - mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=non_transport_error) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): - result = await _db_health_readiness_check() - - assert result["status"] == "disconnected" - mock_prisma.disconnect.assert_not_called() - mock_prisma.connect.assert_not_called() - - -@pytest.mark.asyncio -async def test_db_health_reconnect_disconnect_fails(): - """ - When disconnect() itself raises during the reconnect cycle, - the inner except catches it and returns 'disconnected'. - connect() and the second health_check() are never called. - """ - transport_error = ClientNotConnectedError() - mock_prisma = MagicMock() mock_prisma.health_check = AsyncMock(side_effect=transport_error) - mock_prisma.disconnect = AsyncMock(side_effect=RuntimeError("already closed")) - mock_prisma.connect = AsyncMock() + mock_prisma.attempt_db_reconnect = AsyncMock( + side_effect=RuntimeError("reconnect failed") + ) _health_endpoints_module.db_health_cache = { "status": "connected", "last_updated": datetime.now() - timedelta(seconds=20), } - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): result = await _db_health_readiness_check() assert result["status"] == "disconnected" - mock_prisma.disconnect.assert_called_once() - mock_prisma.connect.assert_not_called() + + +@pytest.mark.asyncio +async def test_db_health_non_transport_error_returns_disconnected(): + """ + When health_check raises a non-transport error (e.g. data-layer error), + is_database_transport_error returns False so reconnect is skipped. + Returns 'disconnected' without raising and without calling attempt_db_reconnect. + """ + non_transport_error = PrismaError("UniqueViolationError") + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=non_transport_error) + mock_prisma.attempt_db_reconnect = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "connected", + "last_updated": datetime.now() - timedelta(seconds=20), + } + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await _db_health_readiness_check() + + assert result["status"] == "disconnected" + mock_prisma.attempt_db_reconnect.assert_not_called() @pytest.mark.asyncio From 3745ba2ea786d28683da9e4507e84f7656e0e3bf Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Mon, 20 Apr 2026 15:44:11 -0700 Subject: [PATCH 06/11] replace retired claude-3-haiku-20240307 with claude-haiku-4-5-20251001 in anthropic messages passthrough test Anthropic retired claude-3-haiku-20240307 on 2026-04-20, causing the test_anthropic_messages_litellm_router_non_streaming_with_logging test to 404. Update the model references in this file to the current pinned haiku version. --- .../test_anthropic_messages_passthrough.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py index ae2d66eb955..84b14f9508b 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py @@ -77,7 +77,7 @@ class TestAnthropicDirectAPI(BaseAnthropicMessagesTest): @property def model_config(self) -> Dict[str, Any]: return { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), } @@ -86,7 +86,7 @@ class TestAnthropicDirectAPI(BaseAnthropicMessagesTest): """ This is the model name that is expected to be in the logging payload """ - return "claude-3-haiku-20240307" + return "claude-haiku-4-5-20251001" class TestAnthropicBedrockAPI(BaseAnthropicMessagesTest): @@ -140,7 +140,7 @@ async def test_anthropic_messages_streaming_with_bad_request(): response = await litellm.anthropic.messages.acreate( messages=[{"role": "user", "content": "hi"}], api_key=os.getenv("ANTHROPIC_API_KEY"), - model="claude-3-haiku-20240307", + model="claude-haiku-4-5-20251001", max_tokens=100, stream=True, ) @@ -168,7 +168,7 @@ async def test_anthropic_messages_router_streaming_with_bad_request(): { "model_name": "claude-special-alias", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -205,7 +205,7 @@ async def test_anthropic_messages_litellm_router_non_streaming(): { "model_name": "claude-special-alias", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -243,7 +243,7 @@ async def test_anthropic_messages_litellm_router_routing_strategy(): { "model_name": "claude-special-alias", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -341,7 +341,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Here's a joke for you!"}], - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -355,7 +355,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): { "model_name": MODEL_GROUP, "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -419,7 +419,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): assert "model_info" in litellm_metadata # Verify other call parameters - assert call_kwargs["model"] == "claude-3-haiku-20240307" + assert call_kwargs["model"] == "claude-haiku-4-5-20251001" assert call_kwargs["messages"] == messages assert call_kwargs["max_tokens"] == 100 assert call_kwargs["metadata"] == {"user_id": "hello"} @@ -459,7 +459,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): { "model_name": MODEL_GROUP, "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -496,7 +496,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): assert test_custom_logger.logged_standard_logging_payload["response"] is not None assert ( test_custom_logger.logged_standard_logging_payload["model"] - == "claude-3-haiku-20240307" + == "claude-haiku-4-5-20251001" ) # check logged usage + spend @@ -543,7 +543,7 @@ async def test_anthropic_messages_with_extra_headers(): "text": "Why did the chicken cross the road? To get to the other side!", } ], - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -556,7 +556,7 @@ async def test_anthropic_messages_with_extra_headers(): response = await litellm.anthropic.messages.acreate( messages=messages, api_key=api_key, - model="claude-3-haiku-20240307", + model="claude-haiku-4-5-20251001", max_tokens=100, client=mock_client, provider_specific_header={ @@ -689,7 +689,7 @@ async def test_anthropic_messages_with_thinking(): "text": "Why did the chicken cross the road? To get to the other side!", } ], - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -702,7 +702,7 @@ async def test_anthropic_messages_with_thinking(): response = await litellm.anthropic.messages.acreate( messages=messages, api_key=api_key, - model="claude-3-haiku-20240307", + model="claude-haiku-4-5-20251001", max_tokens=100, client=mock_client, thinking={"budget_tokens": 100}, @@ -717,7 +717,7 @@ async def test_anthropic_messages_with_thinking(): request_body = json.loads(call_kwargs.get("data", {})) print("REQUEST BODY", request_body) assert request_body["max_tokens"] == 100 - assert request_body["model"] == "claude-3-haiku-20240307" + assert request_body["model"] == "claude-haiku-4-5-20251001" assert request_body["messages"] == messages assert request_body["thinking"] == {"budget_tokens": 100} From 505f6e0522eacbf4e4407b5886890805b3a1fea2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 15:53:12 -0700 Subject: [PATCH 07/11] [Fix] Apply black formatting to fix CI lint failures --- litellm/integrations/prometheus.py | 13 ++++++++----- litellm/integrations/prometheus_helpers/__init__.py | 3 +-- .../integrations/websearch_interception/handler.py | 4 +++- litellm/litellm_core_utils/llm_cost_calc/utils.py | 2 +- .../messages/agentic_streaming_iterator.py | 8 +++++--- litellm/llms/github_copilot/authenticator.py | 4 +--- litellm/passthrough/utils.py | 4 +++- litellm/proxy/auth/auth_checks.py | 11 +++++------ litellm/proxy/health_check.py | 9 ++++----- litellm/proxy/hooks/parallel_request_limiter_v3.py | 6 +++--- .../management_endpoints/internal_user_endpoints.py | 8 ++------ .../management_endpoints/organization_endpoints.py | 5 +---- .../proxy/management_endpoints/team_endpoints.py | 3 +-- litellm/types/integrations/prometheus.py | 2 +- 14 files changed, 39 insertions(+), 43 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1d92a9da073..723b142dfad 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -51,6 +51,7 @@ if TYPE_CHECKING: else: AsyncIOScheduler = Any + class PrometheusLogger(CustomLogger): # Class variables or attributes @@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger): amount: float = 1.0, ) -> None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name=metric_name - ), + supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name), enum_values=enum_values, label_context=label_context, ) @@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger): user_api_key = hash_token(user_api_key) - label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request. + label_context = PrometheusLabelFactoryContext( + enum_values + ) # amortized per request. # increment total LLM requests and spend metric self._increment_top_level_request_and_spend_metrics( @@ -3490,7 +3491,9 @@ def _prometheus_labels_from_context( } if UserAPIKeyLabelNames.END_USER.value in filtered_labels: - filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user() + filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ( + ctx.get_resolved_end_user() + ) for sk, val in ctx._custom_by_sanitized_key.items(): if sk in supported_enum_labels: diff --git a/litellm/integrations/prometheus_helpers/__init__.py b/litellm/integrations/prometheus_helpers/__init__.py index 34f4855863e..784ab524dd5 100644 --- a/litellm/integrations/prometheus_helpers/__init__.py +++ b/litellm/integrations/prometheus_helpers/__init__.py @@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext: self.enum_values = enum_values enum_dict = enum_values.model_dump() self._sanitized_enum: Dict[str, Optional[str]] = { - k: _sanitize_prometheus_label_value(v) - for k, v in enum_dict.items() + k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items() } self._custom_by_sanitized_key: Dict[str, Optional[str]] = {} if enum_values.custom_metadata_labels is not None: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 7b4aa7a3f10..41618c72627 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -847,7 +847,9 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs_for_followup = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) full_model_name = agentic_params.get("model", model) verbose_logger.debug( "WebSearchInterception: Built anthropic request patch " diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 3fd913958da..888999504fe 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -684,7 +684,7 @@ def generic_cost_per_token( # noqa: PLR0915 - cache_creation - image_tokens ) - # Clamp to zero: inconsistent streaming usage + # Clamp to zero: inconsistent streaming usage if text_tokens < 0: text_tokens = 0 prompt_tokens_details["text_tokens"] = text_tokens diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 1f14886ca8e..d0780c82d06 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -216,9 +216,11 @@ class AgenticAnthropicStreamingIterator: return [ - f"{b.get('type')}({b.get('name', '')})" - if b.get("type") == "tool_use" - else b.get("type") + ( + f"{b.get('type')}({b.get('name', '')})" + if b.get("type") == "tool_use" + else b.get("type") + ) for b in rebuilt.get("content", []) ] diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index f4698861edc..9de2987b9f6 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -294,9 +294,7 @@ class Authenticator: access_token_url = os.getenv( "GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL ) - client_id = os.getenv( - "GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID - ) + client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) for attempt in range(max_attempts): try: diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index 5dde13f0078..d39a0dda152 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -79,7 +79,9 @@ class BasePassthroughUtils: for header_name, header_value in request_headers.items(): if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): # Strip the 'x-pass-' prefix and normalize to lowercase - actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() + actual_header_name = header_name[ + len(PASS_THROUGH_HEADER_PREFIX) : + ].lower() if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any( actual_header_name.startswith(p) for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d245ec53ece..e19d04a2609 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3126,9 +3126,7 @@ async def _virtual_key_max_budget_alert_check( alert_email_config: Optional[Dict[str, List[str]]] = ( _merge_budget_alert_email_configs( global_cfg=litellm.default_key_max_budget_alert_emails, - per_key_cfg=(valid_token.metadata or {}).get( - "max_budget_alert_emails" - ), + per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), ) ) @@ -3138,7 +3136,9 @@ async def _virtual_key_max_budget_alert_check( (int(k) for k in alert_email_config if k.isdigit()), default=None, ) - if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0): + if min_pct is None or valid_token.spend < valid_token.max_budget * ( + min_pct / 100.0 + ): return call_info = CallInfo( @@ -3164,8 +3164,7 @@ async def _virtual_key_max_budget_alert_check( else: # Old path: existing single 80% threshold — completely unchanged alert_threshold = ( - valid_token.max_budget - * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE ) if ( diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index e0664703d28..7d67750c78f 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -306,7 +306,9 @@ def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool: return "*" in _deployment_model_string_for_health_check(litellm_params) -def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]: +def _resolve_health_check_max_tokens( + model_info: dict, litellm_params: dict +) -> Optional[int]: """ Pick max_tokens for the health check request. @@ -341,10 +343,7 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> return int(tokens_reasoning) if not is_reasoning and tokens_non_reasoning is not None: return int(tokens_non_reasoning) - if ( - is_reasoning - and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None - ): + if is_reasoning and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None: return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING) if BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5c2b3dfe0ee..f29bbd2d9d5 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1570,9 +1570,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): user_api_key_project_id = standard_logging_metadata.get( "user_api_key_project_id" ) - user_api_key_end_user_id = kwargs.get( - "user" - ) or standard_logging_metadata.get("user_api_key_end_user_id") + user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get( + "user_api_key_end_user_id" + ) model_group = get_model_group_from_litellm_kwargs(kwargs) # Get total tokens from response diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 8474f026111..c6d37ace4fe 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2120,9 +2120,7 @@ async def delete_user( for m in all_target_memberships: if not m.organization_id: continue - target_org_ids_by_user.setdefault(m.user_id, set()).add( - m.organization_id - ) + target_org_ids_by_user.setdefault(m.user_id, set()).add(m.organization_id) # check that all teams passed exist for user_id in data.user_ids: @@ -2141,9 +2139,7 @@ async def delete_user( # Org-admin may only delete users whose entire org membership is # within their admin scope. A target with ANY org outside the # caller's scope (or no org at all) requires PROXY_ADMIN. - if not target_org_ids or not target_org_ids.issubset( - caller_admin_org_ids - ): + if not target_org_ids or not target_org_ids.issubset(caller_admin_org_ids): raise HTTPException( status_code=403, detail={ diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index a6a1af971e5..442fae2a4fa 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1078,10 +1078,7 @@ async def organization_member_update( LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ): - if ( - user_api_key_dict.user_role - != LitellmUserRoles.PROXY_ADMIN.value - ): + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, detail={ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8e21b851857..e4886eb6d15 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1570,8 +1570,7 @@ async def update_team( # noqa: PLR0915 current_org_id = getattr(existing_team_row, "organization_id", None) if ( data.organization_id != current_org_id - and user_api_key_dict.user_role - != LitellmUserRoles.PROXY_ADMIN.value + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): # Is the caller org_admin of the destination org? caller_memberships = ( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 338c5a79ce6..43a287f29bc 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -784,7 +784,7 @@ class UserAPIKeyLabelValues: org_id: Optional[str] = None org_alias: Optional[str] = None - #Added for test compatibility. + # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: """ Match former Pydantic behavior: unknown keys are ignored; ``api_key_hash`` maps to From 74169b114a2053a1db96bc42f6e789341dd2da07 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Mon, 20 Apr 2026 16:04:54 -0700 Subject: [PATCH 08/11] replace retired claude-3-haiku-20240307 with claude-haiku-4-5-20251001 in streaming tests --- tests/local_testing/test_streaming.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 3aed0699603..ecac2cfe40e 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1727,7 +1727,7 @@ def test_openai_chat_completion_complete_response_call(): "model", [ "gpt-3.5-turbo", - "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", "o1", ], ) @@ -2247,7 +2247,7 @@ def streaming_and_function_calling_format_tests(idx, chunk): [ # "gpt-3.5-turbo", # "anthropic.claude-3-sonnet-20240229-v1:0", - "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", ], ) def test_streaming_and_function_calling(model): From eee51a99ad922a047f2e7fb379d3894a3d173248 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Mon, 20 Apr 2026 16:10:45 -0700 Subject: [PATCH 09/11] replace retired claude-3-haiku-20240307 with claude-haiku-4-5-20251001 in local_testing part1 and router fallback tests --- tests/local_testing/test_batch_completions.py | 2 +- tests/local_testing/test_function_call_parsing.py | 2 +- tests/local_testing/test_function_calling.py | 4 ++-- tests/local_testing/test_router_fallbacks.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/local_testing/test_batch_completions.py b/tests/local_testing/test_batch_completions.py index 2125a998f84..95bfe5e6e2b 100644 --- a/tests/local_testing/test_batch_completions.py +++ b/tests/local_testing/test_batch_completions.py @@ -72,7 +72,7 @@ def test_batch_completions_models(): def test_batch_completion_models_all_responses(): try: responses = batch_completion_models_all_responses( - models=["gemini/gemini-2.5-flash-lite", "claude-3-haiku-20240307"], + models=["gemini/gemini-2.5-flash-lite", "claude-haiku-4-5-20251001"], messages=[{"role": "user", "content": "write a poem"}], max_tokens=10, ) diff --git a/tests/local_testing/test_function_call_parsing.py b/tests/local_testing/test_function_call_parsing.py index 0351ce70572..f9582fcc574 100644 --- a/tests/local_testing/test_function_call_parsing.py +++ b/tests/local_testing/test_function_call_parsing.py @@ -142,7 +142,7 @@ def trade(model_name: str) -> List[Trade]: # type: ignore @pytest.mark.parametrize( - "model", ["claude-3-haiku-20240307", "anthropic.claude-3-haiku-20240307-v1:0"] + "model", ["claude-haiku-4-5-20251001", "anthropic.claude-3-haiku-20240307-v1:0"] ) @pytest.mark.flaky(retries=6, delay=10) def test_function_call_parsing(model): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 13adb163d5f..b52805c0664 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -47,7 +47,7 @@ def get_current_weather(location, unit="fahrenheit"): [ "gpt-3.5-turbo-1106", "mistral/mistral-large-latest", - "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", "gemini/gemini-2.5-flash-lite", "anthropic.claude-3-sonnet-20240229-v1:0", ], @@ -275,7 +275,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message "anthropic.claude-3-sonnet-20240229-v1:0", "bedrock", ), - ("claude-3-haiku-20240307", "anthropic"), + ("claude-haiku-4-5-20251001", "anthropic"), ], ) @pytest.mark.parametrize( diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 383ad104577..a14e53adbc4 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -1509,7 +1509,7 @@ def test_router_fallbacks_with_wildcard_model_name(): { "model_name": "claude-3-haiku", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), "mock_response": "Hi this is claude!", }, @@ -1555,7 +1555,7 @@ def test_fallbacks_with_different_messages(): { "model_name": "claude-3-haiku", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, }, From bb6209932312e4e7438f5133fb09fc8989f5bd82 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 16:22:10 -0700 Subject: [PATCH 10/11] [Fix] CI - auth_ui_unit_tests: use Postgres sidecar instead of shared DB Run auth_ui_unit_tests against a per-job cimg/postgres:16.0 sidecar with DATABASE_URL pointing at localhost:5432, matching the pattern used by e2e_ui_testing. Seed the schema via 'litellm --skip_server_startup --use_prisma_db_push' so each run starts on a clean DB with the current schema.prisma. --- .circleci/config.yml | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 3d1e22eebd3..9b976462b13 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -439,7 +439,14 @@ jobs: auth: username: ${DOCKERHUB_USERNAME} password: ${DOCKERHUB_PASSWORD} + - image: cimg/postgres:16.0 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test working_directory: ~/project + environment: + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" steps: - checkout @@ -463,12 +470,14 @@ jobs: paths: - ./.venv key: v2-dependencies-{{ checksum "uv.lock" }}-{{ checksum ".circleci/config.yml" }} + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" - run: - name: Run prisma ./docker/entrypoint.sh + name: Seed DB schema via prisma db push command: | set +e - chmod +x docker/entrypoint.sh - ./docker/entrypoint.sh + uv run --no-sync litellm --skip_server_startup --use_prisma_db_push set -e - run: name: Generate Prisma Client From e7bc316db01b0fb1695b381aff110fea435b4aab Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 20 Apr 2026 16:22:12 -0700 Subject: [PATCH 11/11] Litellm krrish staging 04 20 2026 (#26138) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(router): add auto_router/quality_router for quality-tier routing (#25987) * feat(router): add auto_router/quality_router for quality-tier routing Adds a new auto-router type that routes a request to a model at a target quality tier. The quality tier is inferred by re-using the existing ComplexityRouter's classification, then mapped through an admin-configured complexity_to_quality table. Each candidate model declares its own quality_tier in model_info.litellm_routing_preferences. Resolution strategy: exact tier match, else round up to the next higher tier, else fall back to default_model. Co-Authored-By: Claude Opus 4 (1M context) * feat(quality_router): add capability-based filtering Each deployment can declare a `capabilities: List[str]` field in `model_info.litellm_routing_preferences` (e.g. ["vision", "function_calling"]). Requests can pass `litellm_capabilities` in `request_kwargs` to require specific capabilities — the router will only route to deployments whose declared capabilities are a superset. Resolution still walks tier (exact → round up), but at each tier filters by capability before picking. Falls back to default_model only when it also satisfies the required capabilities; otherwise raises rather than silently routing to a model that lacks a required capability. Co-Authored-By: Claude Opus 4 (1M context) * feat(quality_router): expose routing decision in response headers For transparency, expose the QualityRouter's routing decision in the proxy response headers: x-litellm-quality-router-model → picked model_name (e.g. "haiku-vision") x-litellm-quality-router-tier → resolved quality tier (e.g. "1") x-litellm-quality-router-complexity → ComplexityTier name (e.g. "SIMPLE") Mechanism: the pre-routing hook stashes the decision in request_kwargs["metadata"]["quality_router_decision"]. After the call returns, Router.set_response_headers lifts the decision into response._hidden_params["additional_headers"] alongside the existing x-litellm-model-group / x-litellm-model-id headers. Existing metadata keys (trace_id, user_id, etc.) are preserved. Co-Authored-By: Claude Opus 4 (1M context) * feat(quality_router): replace capabilities with keyword override Drops the capability-based filtering in favor of a keyword-based override for v0: - RoutingPreferences.keywords: List[str] (replaces capabilities) — each deployment can declare substring keywords. - If any declared keyword (case-insensitive) appears in the user message, the router short-circuits the complexity-classification flow and routes to the matching deployment. - Tiebreaker for overlapping keyword matches: quality_tier DESC, then cheapest model_info.input_cost_per_token ASC. Unpriced models lose ties to priced ones. Decision metadata + headers now expose the override: x-litellm-quality-router-via → "keyword" | "quality_tier" x-litellm-quality-router-keyword → matched keyword (only on keyword route) x-litellm-quality-router-complexity → complexity tier (only on tier route) Removes: - request_kwargs["litellm_capabilities"] reading - _model_capabilities, _model_supports_capabilities, _first_capable_model_at_tier, capability filter in _resolve_model_for_quality_tier Co-Authored-By: Claude Opus 4 (1M context) * feat(quality_router): add explicit `order` to RoutingPreferences Adds an explicit priority field to RoutingPreferences for resolving collisions deterministically: RoutingPreferences.order: Optional[int] # lower wins; unset = +inf Used as the PRIMARY tiebreaker in two places: 1. Keyword overlap: when multiple deployments declare the same matching keyword, sort by (order ASC, quality_tier DESC, input_cost_per_token ASC, model_name ASC). Explicit always beats implicit. 2. Tier resolution: when multiple deployments share a quality tier, `_resolve_model_for_quality_tier` picks the one with the lowest order. The tier list is now sorted at index-build time. This lets admins make routing decisions explicit when the natural quality-and-price ordering would pick the wrong model. Co-Authored-By: Claude Opus 4 (1M context) * feat(quality_router): reorder tiebreak to (quality, order, price) Changes the tiebreak ordering so quality_tier always wins first, then explicit `order` is used to break ties within the same tier, then price breaks the rest: 1. quality_tier DESC ← best model wins first 2. order ASC ← explicit priority within a tier 3. input_cost_per_token ASC 4. model_name ASC Previously `order` was the primary key — that meant a tier-2 model with `order=1` would beat a tier-3 model with no `order`, which is the wrong default. Now `order` only resolves collisions among same-tier candidates. Tier resolution (within a single tier) keeps the same key minus quality: (order ASC, cost ASC, name). Test renames + flips: - test_explicit_order_overrides_quality_tier → test_quality_wins_over_explicit_order - new: test_order_breaks_tie_within_same_quality_tier Co-Authored-By: Claude Opus 4 (1M context) * fix(quality_router): resolve Greptile review feedback Addresses four P1 findings from PR review plus test coverage: 1. set_model_list missing quality_routers reset - Hot-reloading the Router would leave stale QualityRouter instances pointing at the old model_list. `set_model_list` now clears `self.quality_routers` alongside the other indices. 2. Round-down fallback before default_model - `_resolve_model_for_quality_tier` now rounds DOWN to the closest lower tier after round-up fails, before falling back to `default_model`. Degrades gracefully rather than jumping straight off-tier. 3. RoutingPreferences validation bypass - `_build_tier_index` now instantiates `RoutingPreferences(**prefs)` so invalid shapes (e.g. non-int quality_tier) raise a clear ValueError instead of silently succeeding. 4. Config-ordering dependency - `_tier_to_models` is now built lazily on first access. Previously, eager construction in `__init__` meant a QualityRouter deployment had to appear AFTER all its referenced models in config.yaml, because `Router._create_deployment` populates `model_list` incrementally. Any `available_models` defined after the router entry would silently be reported as missing. Also adds 6 new tests covering each fix: - test_invalid_quality_tier_type_raises_clear_error - test_router_can_be_instantiated_before_its_targets_exist - test_set_model_list_clears_quality_routers_registry - test_rounds_down_when_no_higher_tier_exists - test_rounds_down_prefers_closest_lower_tier - test_prefers_round_up_over_round_down Co-Authored-By: Claude Opus 4 (1M context) * style: apply black 24.10.0 formatting to pre-existing offenders Unblocks the LiteLLM Linting check for this PR — these 12 files are already failing `black --check` on main (the lint workflow only runs on PRs, so main drifts). No behavior changes; formatting-only. Co-Authored-By: Claude Opus 4.7 (1M context) * Update litellm/router.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: Claude Opus 4 (1M context) Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Support /v1/responses in complexity router (#26137) * feat(proxy): add --reload flag for uvicorn hot reload (dev only) Opt-in CLI flag, off by default, no env var. Only affects the uvicorn run path; gunicorn/hypercorn paths and prod (which doesn't pass the flag) are unaffected. * Feature/add audio support for scaleway (#26110) * feat(scaleway): add SCALEWAY to LlmProviders enum * feat(scaleway): add audio transcription config and dispatch wiring Co-Authored-By: Claude Sonnet 4.6 * test(scaleway): add behavior tests for audio transcription config Co-Authored-By: Claude Sonnet 4.6 * chore(scaleway): advertise audio_transcriptions in endpoint-support JSON * docs(scaleway): document audio transcription support * fix(scaleway): address PR review — plain-text response_format + missing-key fail-fast Co-Authored-By: Claude Sonnet 4.6 * test(scaleway): cover new response paths, drop gettysburg.wav coupling Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 * Prompt Compression - add it to the proxy (#25729) * refactor: new agentic loop event hook simplifies how to create logic for tool based multi llm calls * fix: compress - make it work on anthropic input as well * fix(compress.py): working prompt compression for claude code ensures claude code messages can run through proxy easily * docs: add agentic loop hook guide * docs: add agentic_loop_hook to sidebar * fix: fix multiple arguments error * fix: fix tool call loop for compression on streaming /v1/messages * fix: fix linting errors * fix: fix ci/cd errors * feat(litellm_pre_call_utils.py): use claude code session for litellm session id allows claude code logs to be stitched together, making it easy to know they were all part of the same conversation * fix: suppress incorrect mypy warning rE: module * revert: drop PR's changes to litellm/proxy/_experimental/out/ Restores the 34 HTML files under _experimental/out/ to their pre-PR paths (X/index.html -> X.html). All renames are R100 (content unchanged); no other files are touched. * fix: address greptile review comments on PR #25729 - Skip ``kwargs["tools"] = []`` injection when compression is a no-op — Anthropic Messages rejects empty tool arrays on requests that did not originally declare tools. - Move agentic-loop safety guards (fingerprint cycle / max depth) out of the per-callback try/except so they propagate instead of being swallowed by the generic exception handler. Extracted _check_agentic_loop_safety. - Gate generic ``x--session-id`` capture behind the LITELLM_CAPTURE_VENDOR_SESSION_HEADERS env var (off by default) to preserve backwards compatibility; explicit x-litellm-* headers are unaffected. - Fix monkeypatch target in pre-call-hook test to patch the actual module-level binding (litellm.integrations.compression_interception.handler.compress). - Add regression tests for empty-tools skip and opt-in session capture. Co-Authored-By: Claude Opus 4.6 * revert: drop LITELLM_CAPTURE_VENDOR_SESSION_HEADERS flag Generic x--session-id header capture is a new feature and only runs *after* the explicit x-litellm-trace-id / x-litellm-session-id checks, so it does not change behavior for any existing caller that was already using the LiteLLM headers — no backwards-incompatibility to gate. Co-Authored-By: Claude Opus 4.6 * refactor(compress): replace input_type with CallTypes call_type Drop the bespoke ``CompressionInputType`` literal and use the existing ``litellm.types.utils.CallTypes`` enum instead. ``litellm.compress()`` now takes ``call_type: Union[CallTypes, str]`` (default ``CallTypes.completion``) — no new concept to learn, and the enum is already the way the rest of the codebase talks about request shapes. Supported values: ``completion`` / ``acompletion`` (OpenAI chat-completions shape) and ``anthropic_messages`` (Anthropic structured content blocks). Updated: compress(), the compression_interception handler, tests, docs, and the two eval scripts. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 * Support /v1/responses in complexity router Adds cross-format support to the complexity router via the guardrail translation handler dispatch. Adds get_structured_messages to base translation plus OpenAI chat, Responses, and Anthropic handlers. Auto-router helper _extract_text_from_messages handles tool-call and multimodal messages. Widens async_pre_routing_hook messages type to Dict[str, Any]. Fixes https://github.com/BerriAI/litellm/issues/25134 * chore: apply black formatting * fix: fallback to trying each handler when route inference fails --------- Co-authored-by: Ryan Crabbe Co-authored-by: nhyy244 <106547304+nhyy244@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 * test: cover _is_quality_router_deployment and init_quality_router_deployment * fix: reset auto_routers on set_model_list to prevent hot-reload ValueError * style: apply black formatting to websearch_interception and agentic_streaming_iterator --------- Co-authored-by: yuneng-jiang Co-authored-by: Claude Opus 4 (1M context) Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Ryan Crabbe Co-authored-by: nhyy244 <106547304+nhyy244@users.noreply.github.com> --- litellm/integrations/custom_logger.py | 2 +- .../chat/guardrail_translation/handler.py | 37 +- .../guardrail_translation/base_translation.py | 11 + .../chat/guardrail_translation/handler.py | 22 +- .../guardrail_translation/handler.py | 26 +- litellm/proxy/_new_secret_config.yaml | 20 +- litellm/router.py | 132 ++- .../auto_router/auto_router.py | 28 +- .../complexity_router/complexity_router.py | 128 +- .../quality_router/__init__.py | 21 + .../router_strategy/quality_router/config.py | 74 ++ .../quality_router/quality_router.py | 446 +++++++ litellm/types/router.py | 4 + .../test_openai_guardrail_handler.py | 55 + ...test_openai_responses_guardrail_handler.py | 60 + .../router_strategy/test_auto_router.py | 144 ++- .../router_strategy/test_complexity_router.py | 221 +++- .../router_strategy/test_quality_router.py | 1033 +++++++++++++++++ 18 files changed, 2390 insertions(+), 74 deletions(-) create mode 100644 litellm/router_strategy/quality_router/__init__.py create mode 100644 litellm/router_strategy/quality_router/config.py create mode 100644 litellm/router_strategy/quality_router/quality_router.py create mode 100644 tests/test_litellm/router_strategy/test_quality_router.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 36486747c39..300c311f36d 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -240,7 +240,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self, model: str, request_kwargs: Dict, - messages: Optional[List[Dict[str, str]]] = None, + messages: Optional[List[Dict[str, Any]]] = None, input: Optional[Union[str, List]] = None, specific_deployment: Optional[bool] = False, ) -> Optional[PreRoutingHookResponse]: diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index cb430b06940..2bb82f227bb 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -34,6 +34,7 @@ from litellm.types.llms.anthropic import ( ) from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionRequest, ChatCompletionToolCallChunk, ChatCompletionToolParam, ) @@ -67,6 +68,32 @@ class AnthropicMessagesHandler(BaseTranslation): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: + """Translate Anthropic request to OpenAI chat completion format.""" + ( + chat_completion_compatible_request, + _tool_name_mapping, + ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) + ) + return chat_completion_compatible_request + + def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + """ + Convert Anthropic messages request data to OpenAI-spec structured messages. + + Uses the Anthropic-to-OpenAI adapter to translate message format. + """ + messages = data.get("messages") + if messages is None: + return None + chat_completion_compatible_request = self._translate_to_openai(data) + result = cast( + List[AllMessageValues], + chat_completion_compatible_request.get("messages", []), + ) + return result if result else None + async def process_input_messages( self, data: dict, @@ -82,13 +109,7 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply) - ( - chat_completion_compatible_request, - _tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). - anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) - ) + chat_completion_compatible_request = self._translate_to_openai(data) structured_messages = cast( List[AllMessageValues], @@ -103,8 +124,6 @@ class AnthropicMessagesHandler(BaseTranslation): chat_completion_compatible_request.get("tools", []) ) task_mappings: List[Tuple[int, Optional[int]]] = [] - # Track (message_index, content_index) for each text - # content_index is None for string content, int for list content # Step 1: Extract all text content and images for msg_idx, message in enumerate(messages): diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index e1da0dfa29e..1efeb159a3e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -5,6 +5,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import AllMessageValues class BaseTranslation(ABC): @@ -101,6 +102,16 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]: + """ + Convert request data to OpenAI-spec structured messages. + + Override in subclasses for format-specific conversion. + + Returns None if no convertible content is found. + """ + return None + def extract_request_tool_names(self, data: dict) -> List[str]: """ Extract tool names from the request body for allowlist/policy checks. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 2db19dea0b9..86ca6625629 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -48,6 +48,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + """ + Convert chat completions request data to OpenAI-spec structured messages. + + Messages are already in OpenAI format, so this is a simple extraction. + """ + messages = data.get("messages") + if messages is None: + return None + return cast(List[AllMessageValues], messages) + async def process_input_messages( self, data: dict, @@ -68,9 +79,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls_to_check: List[ChatCompletionToolParam] = [] text_task_mappings: List[Tuple[int, Optional[int]]] = [] tool_call_task_mappings: List[Tuple[int, int]] = [] - # text_task_mappings: Track (message_index, content_index) for each text - # content_index is None for string content, int for list content - # tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call # Step 1: Extract all text content, images, and tool calls for msg_idx, message in enumerate(messages): @@ -92,12 +100,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore - if messages: - msg_list = cast(List[AllMessageValues], messages) + structured_messages = self.get_structured_messages(data) + if structured_messages: inputs["structured_messages"] = ( - openai_messages_without_system(msg_list) + openai_messages_without_system(structured_messages) if skip_system - else msg_list + else structured_messages ) # Pass tools (function definitions) to the guardrail tools = data.get("tools") diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 76f40eed71f..f7dd68aec55 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -43,6 +43,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ( + AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam, ) @@ -70,6 +71,24 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + """ + Convert Responses API request data to OpenAI-spec structured messages. + + Transforms `input` (string or ResponseInputParam) and optional + `instructions` into chat completion messages. + """ + input_data = data.get("input") + if input_data is None: + return None + messages = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input_data, + responses_api_request=data, + ) + ) + return cast(List[AllMessageValues], messages) if messages else None + async def process_input_messages( self, data: dict, @@ -86,12 +105,7 @@ class OpenAIResponsesHandler(BaseTranslation): if input_data is None: return data - structured_messages = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input_data, - responses_api_request=data, - ) - ) + structured_messages = self.get_structured_messages(data) # Handle simple string input if isinstance(input_data, str): diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 36c90c28559..427ec46740b 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -30,13 +30,13 @@ model_list: id: claude-sonnet-4-custom-pricing input_cost_per_token: 0.0003 # 100x standard ($0.000003) output_cost_per_token: 0.0015 # 100x standard ($0.000015) - -litellm_settings: - callbacks: ["compression_interception"] - compression_interception_params: - enabled: true - compression_trigger: 100000 -# # optional: -# # embedding_model: "text-embedding-3-small" -# # embedding_model_params: -# # dimensions: 512 \ No newline at end of file + - model_name: my-auto + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + tiers: + SIMPLE: "gpt-4.1-mini" + COMPLEX: claude-sonnet-4-6 + tier_boundaries: + simple_medium: 0.30 + complexity_router_default_model: small-model diff --git a/litellm/router.py b/litellm/router.py index 6572d96f7b9..f6976109bc7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -200,12 +200,16 @@ if TYPE_CHECKING: from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, ) + from litellm.router_strategy.quality_router.quality_router import ( + QualityRouter, + ) Span = Union[_Span, Any] else: Span = Any AutoRouter = Any ComplexityRouter = Any + QualityRouter = Any PreRoutingHookResponse = Any @@ -464,6 +468,7 @@ class Router: ) # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} self.complexity_routers: Dict[str, "ComplexityRouter"] = {} + self.quality_routers: Dict[str, "QualityRouter"] = {} # Initialize model_group_alias early since it's used in set_model_list self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = ( @@ -5884,7 +5889,7 @@ class Router: response = await response ## PROCESS RESPONSE HEADERS response = await self.set_response_headers( - response=response, model_group=model_group + response=response, model_group=model_group, request_kwargs=kwargs ) return response @@ -6814,6 +6819,8 @@ class Router: """ if litellm_params.model.startswith("auto_router/complexity_router"): return False # This is handled by complexity_router + if litellm_params.model.startswith("auto_router/quality_router"): + return False # This is handled by quality_router if litellm_params.model.startswith("auto_router/"): return True return False @@ -6920,6 +6927,58 @@ class Router: ) self.complexity_routers[deployment.model_name] = complexity_router + def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: + """ + Check if the deployment is a quality-router deployment. + + Returns True if the litellm_params model starts with "auto_router/quality_router". + """ + if litellm_params.model.startswith("auto_router/quality_router"): + return True + return False + + def init_quality_router_deployment(self, deployment: Deployment): + """ + Initialize the quality-router deployment. + + Resolves the default model from either `quality_router_default_model` or + `quality_router_config["default_model"]`, then instantiates the + QualityRouter and stores it in `self.quality_routers`. + """ + # Import here to mirror the AutoRouter / ComplexityRouter init pattern + # and avoid circular imports. + from litellm.router_strategy.quality_router.quality_router import ( + QualityRouter, + ) + + quality_router_config: Optional[dict] = ( + deployment.litellm_params.quality_router_config + ) + + default_model: Optional[str] = ( + deployment.litellm_params.quality_router_default_model + ) + if default_model is None and quality_router_config: + default_model = quality_router_config.get("default_model") + + if default_model is None: + raise ValueError( + "quality_router_default_model is required for quality-router deployments, " + "or set default_model in quality_router_config. Please configure it in the litellm_params" + ) + + quality_router: QualityRouter = QualityRouter( + model_name=deployment.model_name, + default_model=default_model, + litellm_router_instance=self, + quality_router_config=quality_router_config, + ) + if deployment.model_name in self.quality_routers: + raise ValueError( + f"Quality-router deployment {deployment.model_name} already exists. Please use a different model name." + ) + self.quality_routers[deployment.model_name] = quality_router + def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: """ Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments @@ -6966,6 +7025,11 @@ class Router: self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index self.team_model_to_deployment_indices = {} # Reset the team_model index + # Reset per-strategy router registries so hot-reload doesn't leave + # stale routers pointing at the old model_list. + self.quality_routers = {} + self.complexity_routers = {} + self.auto_routers = {} self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -7140,6 +7204,12 @@ class Router: ): self.init_complexity_router_deployment(deployment=deployment) + ######################################################### + # Check if this is a quality-router deployment + ######################################################### + if self._is_quality_router_deployment(litellm_params=deployment.litellm_params): + self.init_quality_router_deployment(deployment=deployment) + return deployment def _initialize_deployment_for_pass_through( @@ -8143,7 +8213,10 @@ class Router: return returned_dict async def set_response_headers( - self, response: Any, model_group: Optional[str] = None + self, + response: Any, + model_group: Optional[str] = None, + request_kwargs: Optional[dict] = None, ) -> Any: """ Add the most accurate rate limit headers for a given model response. @@ -8164,6 +8237,45 @@ class Router: additional_headers = response._hidden_params["additional_headers"] # type: ignore + # Lift QualityRouter routing decision into response headers for + # transparency. The decision is stashed in request_kwargs.metadata + # by QualityRouter.async_pre_routing_hook. + metadata = ( + (request_kwargs.get("metadata") or {}) + if isinstance(request_kwargs, dict) + else {} + ) + decision = ( + metadata.get("quality_router_decision") + if isinstance(metadata, dict) + else None + ) + if isinstance(decision, dict): + # Only emit headers for fields that have a meaningful value. + # `complexity_tier` and `matched_keyword` are mutually exclusive + # (the keyword path short-circuits classification), so each + # request emits one or the other but not both. + if decision.get("routed_model") is not None: + additional_headers["x-litellm-quality-router-model"] = str( + decision["routed_model"] + ) + if decision.get("quality_tier") is not None: + additional_headers["x-litellm-quality-router-tier"] = str( + decision["quality_tier"] + ) + if decision.get("routed_via") is not None: + additional_headers["x-litellm-quality-router-via"] = str( + decision["routed_via"] + ) + if decision.get("matched_keyword") is not None: + additional_headers["x-litellm-quality-router-keyword"] = str( + decision["matched_keyword"] + ) + if decision.get("complexity_tier") is not None: + additional_headers["x-litellm-quality-router-complexity"] = str( + decision["complexity_tier"] + ) + if ( "x-ratelimit-remaining-tokens" not in additional_headers and "x-ratelimit-remaining-requests" not in additional_headers @@ -8708,8 +8820,6 @@ class Router: and self.routing_strategy == "latency-based-routing" ): _settings_to_return[var] = self.lowestlatency_logger.routing_args.json() - elif var == "routing_strategy_args": - _settings_to_return[var] = None return _settings_to_return def update_settings(self, **kwargs): @@ -9620,7 +9730,7 @@ class Router: self, model: str, request_kwargs: Dict, - messages: Optional[List[Dict[str, str]]] = None, + messages: Optional[List[Dict[str, Any]]] = None, input: Optional[Union[str, List]] = None, specific_deployment: Optional[bool] = False, ) -> Optional[PreRoutingHookResponse]: @@ -9653,6 +9763,18 @@ class Router: specific_deployment=specific_deployment, ) + ######################################################### + # Check if any quality-router should be used + ######################################################### + if model in self.quality_routers: + return await self.quality_routers[model].async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) + return None def get_available_deployment( diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 4ead7225abc..58b2c5a3912 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -82,11 +82,34 @@ class AutoRouter(CustomLogger): ) return auto_router_routes + @staticmethod + def _extract_text_from_messages(messages: List[Dict[str, Any]]) -> str: + """ + Extract text content from the last user message for routing. + + Handles tool-call conversations (where the last message may be an + assistant or tool message with non-string content) and multimodal + messages (where content is a list of content blocks). + """ + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content") + if content is None: + return "" + if isinstance(content, list): + return " ".join( + block.get("text", "") + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ) + return str(content) + return "" + async def async_pre_routing_hook( self, model: str, request_kwargs: Dict, - messages: Optional[List[Dict[str, str]]] = None, + messages: Optional[List[Dict[str, Any]]] = None, input: Optional[Union[str, List]] = None, specific_deployment: Optional[bool] = False, ) -> Optional["PreRoutingHookResponse"]: @@ -120,8 +143,7 @@ class AutoRouter(CustomLogger): auto_sync=self.auto_sync_value, ) - user_message: Dict[str, str] = messages[-1] - message_content: str = user_message.get("content", "") + message_content = self._extract_text_from_messages(messages) route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer( text=message_content ) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index e51249b1cb1..aa3bcef6392 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -332,45 +332,68 @@ class ComplexityRouter(CustomLogger): f"No model configured for tier {tier_key} and no default_model set" ) - async def async_pre_routing_hook( + def _resolve_messages( self, - model: str, + messages: Optional[List[Dict[str, Any]]], request_kwargs: Dict, - messages: Optional[List[Dict[str, Any]]] = None, - input: Optional[Union[str, List]] = None, - specific_deployment: Optional[bool] = False, - ) -> Optional["PreRoutingHookResponse"]: + ) -> Optional[List[Dict[str, Any]]]: """ - Pre-routing hook called before the routing decision. + Resolve messages from the request, converting from other formats if needed. - Classifies the request by complexity and returns the appropriate model. - - Args: - model: The original model name requested. - request_kwargs: The request kwargs. - messages: The messages in the request. - input: Optional input for embeddings. - specific_deployment: Whether a specific deployment was requested. - - Returns: - PreRoutingHookResponse with the routed model, or None if no routing needed. + Uses the guardrail translation handler dispatch to convert Responses API + ``input`` (or other non-chat-completions formats) into OpenAI-spec messages. """ - from litellm.types.router import PreRoutingHookResponse + if messages: + return messages - if messages is None or len(messages) == 0: - verbose_router_logger.debug( - "ComplexityRouter: No messages provided, skipping routing" - ) - return None + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) + from litellm.llms import load_guardrail_translation_mappings + from litellm.types.utils import CallTypes - # Extract the last user message and the last system prompt + mappings = load_guardrail_translation_mappings() + call_type: Optional[CallTypes] = None + + # 1. Try route-based inference from proxy metadata + route = request_kwargs.get("litellm_metadata", {}).get( + "user_api_key_request_route" + ) + if route: + call_types_list = get_call_types_for_route(route) + if call_types_list: + for ct in call_types_list: + if ct in mappings: + call_type = ct + break + + # 2. Fallback: try each mapped handler until one produces messages + handlers_to_try: List[Any] = [] + if call_type is not None and call_type in mappings: + handlers_to_try.append(mappings[call_type]()) + else: + handlers_to_try.extend(handler_cls() for handler_cls in mappings.values()) + + for handler in handlers_to_try: + structured = handler.get_structured_messages(request_kwargs) + if structured: + return [ + msg if isinstance(msg, dict) else msg.model_dump() # type: ignore + for msg in structured + ] + return None + + @staticmethod + def _extract_user_message_and_system_prompt( + messages: List[Dict[str, Any]], + ) -> Tuple[Optional[str], Optional[str]]: + """Extract the last user message text and last system prompt from messages.""" user_message: Optional[str] = None system_prompt: Optional[str] = None for msg in reversed(messages): role = msg.get("role", "") content = msg.get("content") or "" - # content may be a list of content parts (e.g. [{"type": "text", "text": "..."}]) if isinstance(content, list): text_parts = [ part.get("text", "") @@ -383,6 +406,52 @@ class ComplexityRouter(CustomLogger): user_message = content elif role == "system" and system_prompt is None: system_prompt = content + if user_message is not None and system_prompt is not None: + break + + return user_message, system_prompt + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Dict, + messages: Optional[List[Dict[str, Any]]] = None, + input: Optional[Union[str, List]] = None, + specific_deployment: Optional[bool] = False, + ) -> Optional["PreRoutingHookResponse"]: + """ + Pre-routing hook called before the routing decision. + + Classifies the request by complexity and returns the appropriate model. + Supports chat completions (messages), Responses API (input), and other + formats via the guardrail translation handler dispatch. + + Args: + model: The original model name requested. + request_kwargs: The request kwargs. + messages: The messages in the request. + input: Optional input for Responses API or embeddings. + specific_deployment: Whether a specific deployment was requested. + + Returns: + PreRoutingHookResponse with the routed model, or None if no routing needed. + """ + from litellm.types.router import PreRoutingHookResponse + + resolved_messages = self._resolve_messages(messages, request_kwargs) + + if not resolved_messages: + verbose_router_logger.debug( + "ComplexityRouter: No messages could be resolved, skipping routing" + ) + return None + + # Determine whether the original request used messages directly + has_original_messages = messages is not None and len(messages) > 0 + + user_message, system_prompt = self._extract_user_message_and_system_prompt( + resolved_messages + ) if user_message is None: verbose_router_logger.debug( @@ -391,13 +460,10 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM), - messages=messages, + messages=messages if has_original_messages else None, ) - # Classify the request tier, score, signals = self.classify(user_message, system_prompt) - - # Get the model for this tier routed_model = self.get_model_for_tier(tier) verbose_router_logger.info( @@ -407,5 +473,5 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=routed_model, - messages=messages, + messages=messages if has_original_messages else None, ) diff --git a/litellm/router_strategy/quality_router/__init__.py b/litellm/router_strategy/quality_router/__init__.py new file mode 100644 index 00000000000..5728943448a --- /dev/null +++ b/litellm/router_strategy/quality_router/__init__.py @@ -0,0 +1,21 @@ +""" +Quality-tier auto-router. + +Re-uses the ComplexityRouter's classification to decide a request's complexity, +then maps that complexity to an admin-configured quality tier and resolves the +target model from each candidate's `model_info.litellm_routing_preferences`. +""" + +from .config import ( + DEFAULT_COMPLEXITY_TO_QUALITY, + QualityRouterConfig, + RoutingPreferences, +) +from .quality_router import QualityRouter + +__all__ = [ + "QualityRouter", + "QualityRouterConfig", + "RoutingPreferences", + "DEFAULT_COMPLEXITY_TO_QUALITY", +] diff --git a/litellm/router_strategy/quality_router/config.py b/litellm/router_strategy/quality_router/config.py new file mode 100644 index 00000000000..125ecd5bb9b --- /dev/null +++ b/litellm/router_strategy/quality_router/config.py @@ -0,0 +1,74 @@ +""" +Configuration models for the QualityRouter. +""" + +from typing import Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + +# Default mapping from ComplexityTier name (string) to quality tier (int). +# Higher tier = higher capability requirement. +DEFAULT_COMPLEXITY_TO_QUALITY: Dict[str, int] = { + "SIMPLE": 1, + "MEDIUM": 2, + "COMPLEX": 3, + "REASONING": 4, +} + + +class QualityRouterConfig(BaseModel): + """Configuration for the QualityRouter.""" + + available_models: List[str] = Field( + default_factory=list, + description=( + "List of candidate model names this router may route to. Each model " + "must declare its quality_tier in model_info.litellm_routing_preferences." + ), + ) + + default_model: Optional[str] = Field( + default=None, + description="Fallback model when no quality tier resolves.", + ) + + complexity_to_quality: Dict[str, int] = Field( + default_factory=lambda: DEFAULT_COMPLEXITY_TO_QUALITY.copy(), + description="Mapping from ComplexityTier name to quality tier (int).", + ) + + model_config = ConfigDict(extra="allow") + + +class RoutingPreferences(BaseModel): + """Per-deployment routing preferences declared on model_info.""" + + quality_tier: int = Field( + ..., + description="The quality tier this deployment satisfies.", + ) + + keywords: List[str] = Field( + default_factory=list, + description=( + "Substring keywords (case-insensitive) that, when present in the " + "user message, route the request to this deployment. See `order` " + "for explicit collision handling, otherwise ties fall through to " + "(highest quality_tier, then cheapest model_info.input_cost_per_token)." + ), + ) + + order: Optional[int] = Field( + default=None, + description=( + "Explicit priority used to break ties between deployments at the " + "same quality tier. Lower values win. Applies both to keyword " + "collisions and to picking between multiple deployments at the " + "same quality_tier. Tiebreak order is " + "(quality_tier DESC, order ASC, input_cost_per_token ASC, " + "model_name ASC) — quality always wins first, then explicit " + "order, then price." + ), + ) + + model_config = ConfigDict(extra="allow") diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py new file mode 100644 index 00000000000..a79b4384f5e --- /dev/null +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -0,0 +1,446 @@ +""" +Quality-tier Auto Router. + +Routes a request to a model at a target quality tier. The quality tier is +inferred by re-using the existing ComplexityRouter's classification, then +mapped through an admin-configured `complexity_to_quality` table. Each +candidate model declares its own `quality_tier` in +`model_info.litellm_routing_preferences`. + +Optional keyword override: deployments may also declare `keywords` in +`litellm_routing_preferences`. If any declared keyword appears in the user +message (case-insensitive substring match), the router short-circuits the +complexity-classification flow and routes to the matching deployment. When +multiple deployments match, ties are broken by (highest quality_tier first, +then cheapest `model_info.input_cost_per_token`). +""" + +import math +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.complexity_router.complexity_router import ( + ComplexityRouter, +) + +from .config import QualityRouterConfig, RoutingPreferences + +if TYPE_CHECKING: + from litellm.router import Router + from litellm.types.router import PreRoutingHookResponse +else: + Router = Any + PreRoutingHookResponse = Any + + +class QualityRouter(CustomLogger): + """ + Routes requests to a model at a target quality tier, with an optional + keyword override. + """ + + def __init__( + self, + model_name: str, + litellm_router_instance: "Router", + default_model: Optional[str] = None, + quality_router_config: Optional[Dict[str, Any]] = None, + ): + self.model_name = model_name + self.litellm_router_instance = litellm_router_instance + + if quality_router_config: + self.config = QualityRouterConfig(**quality_router_config) + else: + self.config = QualityRouterConfig() + + # Explicit default_model arg overrides anything in the config dict. + if default_model: + self.config.default_model = default_model + + # Internal scorer — re-use the existing rule-based classifier. + self._scorer = ComplexityRouter( + model_name=f"{model_name}::scorer", + litellm_router_instance=litellm_router_instance, + ) + + # Per-model indices populated alongside the tier index. `_model_keywords` + # stores keywords lowercased so we can substring-match against the + # lowercased user message in O(total-keyword-count). `_model_quality`, + # `_model_cost`, and `_model_order` drive tiebreaking — `_model_order` + # is the explicit priority (lower wins, unset = +inf). + self._model_keywords: Dict[str, List[str]] = {} + self._model_quality: Dict[str, int] = {} + self._model_cost: Dict[str, Optional[float]] = {} + self._model_order: Dict[str, Optional[int]] = {} + + # Tier → models index. Built lazily on first access so the QualityRouter + # deployment does NOT need to appear after all its referenced models in + # the config — when `_build_tier_index` runs eagerly in `__init__`, the + # router instance's `model_list` is still being assembled incrementally + # by `_create_deployment`, and any `available_models` defined AFTER the + # router entry in config.yaml would silently be reported as missing. + self._tier_to_models_cache: Optional[Dict[int, List[str]]] = None + + verbose_router_logger.debug( + f"QualityRouter initialized for {model_name} with " + f"available_models={self.config.available_models}, " + f"default_model={self.config.default_model}" + ) + + @property + def _tier_to_models(self) -> Dict[int, List[str]]: + """Lazy tier→models index; built on first access.""" + if self._tier_to_models_cache is None: + self._tier_to_models_cache = self._build_tier_index() + return self._tier_to_models_cache + + def _get_routing_preferences(self, deployment: Any) -> Optional[Dict[str, Any]]: + """ + Extract litellm_routing_preferences from a deployment, handling both + dict-shaped and Pydantic-object-shaped deployments. + """ + # Dict-shaped deployment. + if isinstance(deployment, dict): + model_info = deployment.get("model_info") or {} + if isinstance(model_info, dict): + return model_info.get("litellm_routing_preferences") + # Pydantic ModelInfo nested in a dict. + return getattr(model_info, "litellm_routing_preferences", None) + + # Pydantic-object deployment. + model_info = getattr(deployment, "model_info", None) + if model_info is None: + return None + if isinstance(model_info, dict): + return model_info.get("litellm_routing_preferences") + return getattr(model_info, "litellm_routing_preferences", None) + + def _get_deployment_input_cost(self, deployment: Any) -> Optional[float]: + """ + Extract `input_cost_per_token` from a deployment's model_info. + + Returns None when not declared — None is treated as "infinite cost" + for the cheapest-tiebreak ordering, so unpriced models lose ties to + priced ones. (Admins who want a model to win on price must declare it.) + """ + if isinstance(deployment, dict): + model_info = deployment.get("model_info") or {} + else: + model_info = getattr(deployment, "model_info", None) or {} + + if isinstance(model_info, dict): + cost = model_info.get("input_cost_per_token") + else: + cost = getattr(model_info, "input_cost_per_token", None) + + if cost is None: + return None + try: + return float(cost) + except (TypeError, ValueError): + return None + + def _get_deployment_model_name(self, deployment: Any) -> Optional[str]: + """Extract `model_name` from a dict- or object-shaped deployment.""" + if isinstance(deployment, dict): + return deployment.get("model_name") + return getattr(deployment, "model_name", None) + + def _build_tier_index(self) -> Dict[int, List[str]]: + """ + Build {quality_tier: [model_name, ...]} for every model in + `available_models`, plus side indices `_model_keywords`, + `_model_quality`, and `_model_cost`. Raises if any listed model is + missing `litellm_routing_preferences`. + """ + model_list = getattr(self.litellm_router_instance, "model_list", None) or [] + available = set(self.config.available_models) + + # Track which available models we've matched so we can error on missing. + seen: Dict[str, bool] = {name: False for name in available} + tier_to_models: Dict[int, List[str]] = {} + + for deployment in model_list: + name = self._get_deployment_model_name(deployment) + if name is None or name not in available: + continue + + raw_prefs = self._get_routing_preferences(deployment) + if raw_prefs is None: + raise ValueError( + f"QualityRouter: model '{name}' is listed in available_models " + f"but has no model_info.litellm_routing_preferences" + ) + + # Validate via the Pydantic model so we get a clear error for + # missing quality_tier, wrong types, etc. This also means + # `RoutingPreferences` is the single source of truth for the + # accepted shape — readers relied on raw dicts before. + try: + if isinstance(raw_prefs, RoutingPreferences): + prefs = raw_prefs + elif isinstance(raw_prefs, dict): + prefs = RoutingPreferences(**raw_prefs) + else: + # A Pydantic object of some other shape — coerce via its dict. + prefs = RoutingPreferences( + **( + raw_prefs.model_dump() + if hasattr(raw_prefs, "model_dump") + else dict(raw_prefs) + ) + ) + except Exception as e: + raise ValueError( + f"QualityRouter: model '{name}' has invalid " + f"litellm_routing_preferences: {e}" + ) from e + + tier_int = int(prefs.quality_tier) + tier_to_models.setdefault(tier_int, []).append(name) + self._model_keywords[name] = [str(k).lower() for k in prefs.keywords if k] + self._model_quality[name] = tier_int + self._model_cost[name] = self._get_deployment_input_cost(deployment) + self._model_order[name] = prefs.order + seen[name] = True + + missing = [name for name, found in seen.items() if not found] + if missing: + raise ValueError( + f"QualityRouter: the following available_models are not present in " + f"the router's model_list (or are missing routing preferences): {missing}" + ) + + # Sort each tier's model list so `_resolve_model_for_quality_tier` + # (which picks index [0]) honors (order ASC, cost ASC, name ASC). + # Quality is moot within a single tier; keep parity with the keyword + # tiebreak by ordering on (order, cost, name) here. + for models in tier_to_models.values(): + models.sort(key=lambda n: (self._order_key(n), self._cost_key(n), n)) + + return tier_to_models + + def _order_key(self, model_name: str) -> float: + """`order` lookup as a float — unset becomes +inf so explicit wins.""" + order = self._model_order.get(model_name) + return float(order) if order is not None else math.inf + + def _cost_key(self, model_name: str) -> float: + """`input_cost_per_token` as a float — unset becomes +inf.""" + cost = self._model_cost.get(model_name) + return float(cost) if cost is not None else math.inf + + def _keyword_override(self, user_message: str) -> Optional[Tuple[str, str]]: + """ + Find a deployment whose declared keywords appear in `user_message`. + + Returns (model_name, matched_keyword) or None when no keyword matches. + When multiple deployments match, sorts by: + 1. quality_tier DESC (best quality always wins first) + 2. `order` ASC (explicit priority — unset = +inf so explicit wins + within the same tier) + 3. input_cost_per_token ASC (unpriced = +inf so priced wins) + 4. model_name ASC (deterministic stability) + """ + # Touch the lazy index so `_model_keywords` / `_model_quality` / + # `_model_cost` / `_model_order` are populated. + _ = self._tier_to_models + + text = user_message.lower() + + matches: List[Tuple[str, str]] = [] # (model_name, matched_keyword) + for model_name, keywords in self._model_keywords.items(): + for kw in keywords: + if kw and kw in text: + matches.append((model_name, kw)) + break # one match per model is enough + + if not matches: + return None + + def sort_key(match: Tuple[str, str]) -> Tuple[int, float, float, str]: + name = match[0] + quality = self._model_quality.get(name, 0) + order_val = self._order_key(name) + cost = self._model_cost.get(name) + cost_val = cost if cost is not None else math.inf + # Negate quality so higher tier sorts first under ASC sort. + return (-quality, order_val, cost_val, name) + + matches.sort(key=sort_key) + return matches[0] + + def _resolve_model_for_quality_tier(self, tier: int) -> str: + """ + Resolve a quality tier to a concrete model name. + + Strategy: + 1. Exact tier match → first model registered at that tier. + 2. Round UP to the next higher tier that has a model (closer to a + request we might lack capacity for). + 3. Round DOWN to the closest lower tier that has a model (degrade + gracefully instead of jumping straight to `default_model`, + which may be off-tier). + 4. Fall back to `config.default_model`. + 5. Otherwise raise. + """ + tier_index = self._tier_to_models + if tier in tier_index and tier_index[tier]: + return tier_index[tier][0] + + # Round up. + higher_tiers = sorted(t for t in tier_index if t > tier) + for t in higher_tiers: + if tier_index[t]: + return tier_index[t][0] + + # Round down — closest lower tier first. + lower_tiers = sorted((t for t in tier_index if t < tier), reverse=True) + for t in lower_tiers: + if tier_index[t]: + return tier_index[t][0] + + if self.config.default_model: + return self.config.default_model + + raise ValueError( + f"QualityRouter: no model available for quality tier {tier} and " + f"no default_model configured" + ) + + def _stash_decision( + self, + request_kwargs: Optional[Dict[str, Any]], + decision: Dict[str, Any], + ) -> None: + """ + Stash the routing decision in request_kwargs.metadata so the Router can + lift it into response headers (`x-litellm-quality-router-*`). The same + dict object flows from here through to `make_call.set_response_headers`. + """ + if request_kwargs is None: + return + metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(metadata, dict): + metadata["quality_router_decision"] = decision + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Dict, + messages: Optional[List[Dict[str, Any]]] = None, + input: Optional[Union[str, List]] = None, + specific_deployment: Optional[bool] = False, + ) -> Optional["PreRoutingHookResponse"]: + """Try keyword override first; fall back to complexity-tier routing.""" + from litellm.types.router import PreRoutingHookResponse + + if messages is None or len(messages) == 0: + verbose_router_logger.debug( + "QualityRouter: No messages provided, skipping routing" + ) + return None + + # Extract last user message and last system prompt — same rules as + # ComplexityRouter.async_pre_routing_hook. + user_message: Optional[str] = None + system_prompt: Optional[str] = None + + for msg in reversed(messages): + role = msg.get("role", "") + content = msg.get("content") or "" + if isinstance(content, list): + text_parts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + content = " ".join(text_parts).strip() + if isinstance(content, str) and content: + if role == "user" and user_message is None: + user_message = content + elif role == "system" and system_prompt is None: + system_prompt = content + + if user_message is None: + verbose_router_logger.debug( + "QualityRouter: No user message found, routing to default model" + ) + if not self.config.default_model: + raise ValueError( + "QualityRouter: no user message and no default_model configured" + ) + return PreRoutingHookResponse( + model=self.config.default_model, + messages=messages, + ) + + # Try keyword override first — it short-circuits complexity classification. + keyword_match = self._keyword_override(user_message) + if keyword_match is not None: + routed_model, matched_keyword = keyword_match + verbose_router_logger.info( + f"QualityRouter: keyword override matched='{matched_keyword}' " + f"routed_model={routed_model} " + f"(quality_tier={self._model_quality.get(routed_model)}, " + f"input_cost_per_token={self._model_cost.get(routed_model)})" + ) + self._stash_decision( + request_kwargs, + { + "router_model_name": self.model_name, + "routed_model": routed_model, + "routed_via": "keyword", + "matched_keyword": matched_keyword, + "quality_tier": self._model_quality.get(routed_model), + "complexity_tier": None, + }, + ) + return PreRoutingHookResponse( + model=routed_model, + messages=messages, + ) + + # No keyword match → complexity classification flow. + complexity_tier, score, signals = self._scorer.classify( + user_message, system_prompt + ) + complexity_name = ( + complexity_tier.value + if hasattr(complexity_tier, "value") + else str(complexity_tier) + ) + + quality_tier = self.config.complexity_to_quality.get(complexity_name) + if quality_tier is None: + raise ValueError( + f"QualityRouter: complexity tier '{complexity_name}' not present " + f"in complexity_to_quality mapping {self.config.complexity_to_quality}" + ) + + routed_model = self._resolve_model_for_quality_tier(int(quality_tier)) + + verbose_router_logger.info( + f"QualityRouter: complexity={complexity_name}, score={score:.3f}, " + f"signals={signals}, quality_tier={quality_tier}, " + f"routed_model={routed_model}" + ) + + self._stash_decision( + request_kwargs, + { + "router_model_name": self.model_name, + "routed_model": routed_model, + "routed_via": "quality_tier", + "matched_keyword": None, + "quality_tier": int(quality_tier), + "complexity_tier": complexity_name, + }, + ) + + return PreRoutingHookResponse( + model=routed_model, + messages=messages, + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 6bd64915d79..fb71e1f6491 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -221,6 +221,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): complexity_router_config: Optional[Dict] = None complexity_router_default_model: Optional[str] = None + # quality-router params + quality_router_config: Optional[Dict] = None + quality_router_default_model: Optional[str] = None + # Batch/File API Params s3_bucket_name: Optional[str] = None s3_encryption_key_id: Optional[str] = None diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a4ac4c94d29..a2c37002942 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -891,6 +891,61 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert result == responses_so_far +class TestGetStructuredMessages: + """Test the get_structured_messages method.""" + + def test_should_return_messages_from_chat_completions_request(self): + """Test that messages are returned from a chat completions request.""" + handler = OpenAIChatCompletionsHandler() + data = { + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + } + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[1]["role"] == "user" + + def test_should_return_none_when_no_messages(self): + """Test that None is returned when no messages key exists.""" + handler = OpenAIChatCompletionsHandler() + data = {"model": "gpt-4"} + result = handler.get_structured_messages(data) + assert result is None + + def test_should_return_none_for_none_messages(self): + """Test that None is returned when messages is explicitly None.""" + handler = OpenAIChatCompletionsHandler() + data = {"messages": None} + result = handler.get_structured_messages(data) + assert result is None + + def test_should_handle_multimodal_content(self): + """Test that messages with multimodal content are returned.""" + handler = OpenAIChatCompletionsHandler() + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + }, + ], + } + ] + } + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) == 1 + assert isinstance(result[0]["content"], list) + + if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index ccece8018ff..aee6ccc2e76 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -995,3 +995,63 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: # Should return the responses assert result == responses_so_far + + +class TestGetStructuredMessages: + """Test the get_structured_messages method for Responses API handler.""" + + def test_should_convert_string_input_to_messages(self): + """Test that a simple string input is converted to OpenAI messages.""" + handler = OpenAIResponsesHandler() + data = {"input": "What is the capital of France?"} + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) >= 1 + found_user = False + for msg in result: + if isinstance(msg, dict) and msg.get("role") == "user": + found_user = True + break + assert found_user, f"Expected a user message, got: {result}" + + def test_should_convert_list_input_to_messages(self): + """Test that list input (ResponseInputParam) is converted to OpenAI messages.""" + handler = OpenAIResponsesHandler() + data = { + "input": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + {"role": "user", "content": "How are you?"}, + ] + } + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) >= 3 + + def test_should_include_instructions_as_system_message(self): + """Test that instructions are included as a system message.""" + handler = OpenAIResponsesHandler() + data = { + "input": "Roll a d20", + "instructions": "You are a helpful dungeon master.", + } + result = handler.get_structured_messages(data) + assert result is not None + has_system = any( + isinstance(msg, dict) and msg.get("role") == "system" for msg in result + ) + assert has_system, f"Expected system message from instructions, got: {result}" + + def test_should_return_none_when_no_input(self): + """Test that None is returned when input key is missing.""" + handler = OpenAIResponsesHandler() + data = {"model": "gpt-4o"} + result = handler.get_structured_messages(data) + assert result is None + + def test_should_return_none_for_none_input(self): + """Test that None is returned when input is explicitly None.""" + handler = OpenAIResponsesHandler() + data = {"input": None} + result = handler.get_structured_messages(data) + assert result is None diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index caff2bc8f10..cb46a4ae553 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -12,7 +12,148 @@ sys.path.insert( from litellm.router_strategy.auto_router.auto_router import AutoRouter -pytestmark = pytest.mark.skip(reason="Skipping auto router tests - beta feature") +pytestmark_skip_beta = pytest.mark.skip( + reason="Skipping auto router tests - beta feature" +) + + +class TestExtractTextFromMessages: + """Tests for AutoRouter._extract_text_from_messages (no semantic_router dependency).""" + + def test_should_extract_content_from_simple_user_message(self): + messages = [{"role": "user", "content": "Hello world"}] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "Hello world" + + def test_should_extract_last_user_message_from_tool_call_conversation(self): + messages = [ + {"role": "user", "content": "What's the weather in NYC?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": "72°F and sunny", + }, + {"role": "user", "content": "Now tell me about London"}, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "Now tell me about London" + + def test_should_find_user_message_when_last_message_is_assistant_with_tool_calls( + self, + ): + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "What's the weather?" + + def test_should_find_user_message_when_last_message_is_tool_response(self): + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc", + "content": "72°F and sunny", + }, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "What's the weather?" + + def test_should_handle_multimodal_content_list(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.png"}, + }, + ], + } + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "What's in this image?" + + def test_should_handle_multimodal_content_with_multiple_text_blocks(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "First part"}, + {"type": "text", "text": "Second part"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.png"}, + }, + ], + } + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "First part Second part" + + def test_should_return_empty_string_when_user_content_is_none(self): + messages = [{"role": "user", "content": None}] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "" + + def test_should_return_empty_string_when_no_user_messages(self): + messages = [ + {"role": "system", "content": "You are a helpful assistant"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "" + + def test_should_return_empty_string_for_empty_messages_list(self): + result = AutoRouter._extract_text_from_messages([]) + assert result == "" @pytest.fixture @@ -41,6 +182,7 @@ def mock_route_choice(): return mock_choice +@pytestmark_skip_beta class TestAutoRouter: """Test class for AutoRouter methods.""" diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8d36fc2ba32..e68ea863d82 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -7,7 +7,7 @@ Tests the rule-based complexity scoring and tier assignment logic. import os import sys from typing import Dict, List -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -828,3 +828,222 @@ class TestRouterComplexityDeploymentMethods: ) router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + + +class TestAsyncPreRoutingHookMultiFormat: + """Test async_pre_routing_hook with multiple input formats.""" + + @pytest.mark.asyncio + async def test_should_route_with_chat_completions_messages(self, complexity_router): + """Test routing with standard chat completions messages.""" + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "What is 2+2?"}], + ) + assert result is not None + assert result.model is not None + assert result.messages is not None + + @pytest.mark.asyncio + async def test_should_route_with_responses_api_string_input( + self, complexity_router + ): + """Test routing with Responses API string input via handler dispatch.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"input": "What is the capital of France?"}, + messages=None, + input="What is the capital of France?", + ) + + assert result is not None + assert result.model is not None + # messages should be None since the original request didn't have messages + assert result.messages is None + + @pytest.mark.asyncio + async def test_should_route_with_responses_api_list_input(self, complexity_router): + """Test routing with Responses API list input via handler dispatch.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + list_input = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + { + "role": "user", + "content": "Write a Python function to sort a list using merge sort", + }, + ] + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"input": list_input}, + messages=None, + input=list_input, + ) + + assert result is not None + assert result.model is not None + assert result.messages is None + + @pytest.mark.asyncio + async def test_should_use_route_based_inference(self, complexity_router): + """Test that route-based call type inference is used when available.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={ + "input": "Roll 2d4+1", + "litellm_metadata": { + "user_api_key_request_route": "/v1/responses", + }, + }, + messages=None, + ) + + assert result is not None + assert result.model is not None + + @pytest.mark.asyncio + async def test_should_return_none_when_no_messages_or_input( + self, complexity_router + ): + """Test that None is returned when neither messages nor input is available.""" + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=None, + input=None, + ) + assert result is None + + @pytest.mark.asyncio + async def test_should_prefer_original_messages_over_conversion( + self, complexity_router + ): + """Test that original messages are used when both messages and input are available.""" + messages = [{"role": "user", "content": "What is 2+2?"}] + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"input": "This should be ignored"}, + messages=messages, + ) + assert result is not None + assert result.messages == messages + + @pytest.mark.asyncio + async def test_should_include_instructions_in_classification( + self, complexity_router + ): + """Test that Responses API instructions influence classification via system message.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={ + "input": "Write merge sort", + "instructions": "You are an expert Python developer. Use advanced algorithms and optimize for performance.", + }, + messages=None, + ) + + assert result is not None + assert result.model is not None + + +class TestExtractUserMessageAndSystemPrompt: + """Test the _extract_user_message_and_system_prompt static method.""" + + def test_should_extract_user_message(self): + """Test extraction of the last user message.""" + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + {"role": "user", "content": "How are you?"}, + ] + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + messages + ) + assert user_msg == "How are you?" + assert sys_prompt == "You are helpful." + + def test_should_handle_no_user_message(self): + """Test when there is no user message.""" + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "assistant", "content": "Hi!"}, + ] + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + messages + ) + assert user_msg is None + assert sys_prompt == "You are helpful." + + def test_should_handle_multipart_content(self): + """Test extraction from multipart content messages.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.png"}, + }, + ], + } + ] + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + messages + ) + assert user_msg == "Describe this image" + assert sys_prompt is None + + def test_should_handle_empty_messages(self): + """Test with empty messages list.""" + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + [] + ) + assert user_msg is None + assert sys_prompt is None diff --git a/tests/test_litellm/router_strategy/test_quality_router.py b/tests/test_litellm/router_strategy/test_quality_router.py new file mode 100644 index 00000000000..01574cb980d --- /dev/null +++ b/tests/test_litellm/router_strategy/test_quality_router.py @@ -0,0 +1,1033 @@ +""" +Tests for the QualityRouter. + +Covers: +- Tier index construction from `model_info.litellm_routing_preferences`. +- Quality-tier resolution (exact, round-up, default fallback). +- Keyword override (match, tiebreaking by quality + price). +- Pre-routing hook end-to-end. +- Decision metadata stash + Router.set_response_headers lift. +""" + +import os +import sys +from typing import Any, Dict, List +from unittest.mock import MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.router_strategy.quality_router.config import ( + DEFAULT_COMPLEXITY_TO_QUALITY, +) +from litellm.router_strategy.quality_router.quality_router import QualityRouter + + +def _make_model_list(spec: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Build a router model_list from a compact spec. + + spec entry shape: { + "model_name": str, + "quality_tier": Optional[int], + "keywords": Optional[List[str]], + "order": Optional[int], + "input_cost_per_token": Optional[float], + } + If quality_tier is None, the deployment is created without + `litellm_routing_preferences`. + """ + out: List[Dict[str, Any]] = [] + for entry in spec: + model_info: Dict[str, Any] = {"id": f"id-{entry['model_name']}"} + if entry.get("quality_tier") is not None: + prefs: Dict[str, Any] = {"quality_tier": entry["quality_tier"]} + if "keywords" in entry: + prefs["keywords"] = entry["keywords"] + if "order" in entry: + prefs["order"] = entry["order"] + model_info["litellm_routing_preferences"] = prefs + if "input_cost_per_token" in entry: + model_info["input_cost_per_token"] = entry["input_cost_per_token"] + out.append( + { + "model_name": entry["model_name"], + "litellm_params": {"model": f"openai/{entry['model_name']}"}, + "model_info": model_info, + } + ) + return out + + +@pytest.fixture +def four_tier_model_list() -> List[Dict[str, Any]]: + """A standard haiku(1)/sonnet(2)/opus(3)/opus-next(4) model list.""" + return _make_model_list( + [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": 2}, + {"model_name": "opus", "quality_tier": 3}, + {"model_name": "opus-next", "quality_tier": 4}, + ] + ) + + +@pytest.fixture +def mock_router(four_tier_model_list): + """A MagicMock router preloaded with the four-tier model list.""" + router = MagicMock() + router.model_list = four_tier_model_list + return router + + +@pytest.fixture +def quality_router(mock_router) -> QualityRouter: + """Default QualityRouter wired to all four tiers.""" + config = { + "available_models": ["haiku", "sonnet", "opus", "opus-next"], + "complexity_to_quality": DEFAULT_COMPLEXITY_TO_QUALITY, + } + return QualityRouter( + model_name="quality-router-test", + litellm_router_instance=mock_router, + default_model="haiku", + quality_router_config=config, + ) + + +# ─── Tier index ───────────────────────────────────────────────────────────── + + +class TestTierIndex: + def test_builds_correct_tier_to_models_map(self, quality_router): + assert quality_router._tier_to_models == { + 1: ["haiku"], + 2: ["sonnet"], + 3: ["opus"], + 4: ["opus-next"], + } + + def test_ignores_models_not_in_available_models(self, four_tier_model_list): + # Add a model the config doesn't list — it should be ignored. + extra = _make_model_list([{"model_name": "ghost", "quality_tier": 5}]) + router = MagicMock() + router.model_list = four_tier_model_list + extra + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={ + "available_models": ["haiku", "sonnet", "opus", "opus-next"] + }, + ) + + for models in qr._tier_to_models.values(): + assert "ghost" not in models + + def test_raises_when_routing_preferences_missing(self): + # `sonnet` is in available_models but has no preferences. + ml = _make_model_list( + [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": None}, + ] + ) + router = MagicMock() + router.model_list = ml + + # Construction succeeds (tier index is lazy); the error surfaces on + # first use so the router entry doesn't have to appear after all of + # its referenced models in config.yaml. + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku", "sonnet"]}, + ) + with pytest.raises(ValueError, match="sonnet"): + _ = qr._tier_to_models + + +# ─── Resolve model for quality tier ───────────────────────────────────────── + + +class TestResolveModelForQualityTier: + def test_exact_match(self, quality_router): + assert quality_router._resolve_model_for_quality_tier(2) == "sonnet" + assert quality_router._resolve_model_for_quality_tier(4) == "opus-next" + + def test_rounds_up_when_tier_missing(self, mock_router): + # Available tiers: 1, 3, 4. Asking for 2 should round up to 3. + spec = [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "opus", "quality_tier": 3}, + {"model_name": "opus-next", "quality_tier": 4}, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku", "opus", "opus-next"]}, + ) + + assert qr._resolve_model_for_quality_tier(2) == "opus" + + def test_rounds_down_when_no_higher_tier_exists(self): + # Only tier 1 available. Asking for tier 4 rounds up (nothing), then + # rounds DOWN to the closest lower tier — tier 1. + spec = [{"model_name": "haiku", "quality_tier": 1}] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="emergency-default", + quality_router_config={"available_models": ["haiku"]}, + ) + + assert qr._resolve_model_for_quality_tier(4) == "haiku" + + def test_rounds_down_prefers_closest_lower_tier(self): + # Available: 1, 2. Asking for 4 rounds down to tier 2 (not tier 1). + spec = [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": 2}, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="emergency-default", + quality_router_config={"available_models": ["haiku", "sonnet"]}, + ) + + assert qr._resolve_model_for_quality_tier(4) == "sonnet" + + def test_prefers_round_up_over_round_down(self): + # Available: 1, 3. Asking for 2 rounds UP to 3, not DOWN to 1. + spec = [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "opus", "quality_tier": 3}, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="emergency-default", + quality_router_config={"available_models": ["haiku", "opus"]}, + ) + + assert qr._resolve_model_for_quality_tier(2) == "opus" + + +# ─── RoutingPreferences validation ───────────────────────────────────────── + + +class TestRoutingPreferencesValidation: + def test_invalid_quality_tier_type_raises_clear_error(self): + # quality_tier must be an int — pass a non-coercible string. + ml = [ + { + "model_name": "haiku", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": { + "id": "id-haiku", + "litellm_routing_preferences": {"quality_tier": "not-an-int"}, + }, + } + ] + router = MagicMock() + router.model_list = ml + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku"]}, + ) + with pytest.raises(ValueError, match="invalid litellm_routing_preferences"): + _ = qr._tier_to_models + + +# ─── Config-ordering independence (lazy index build) ─────────────────────── + + +class TestConfigOrderingIndependence: + def test_router_can_be_instantiated_before_its_targets_exist(self): + # Build a router instance whose referenced model_list is EMPTY at + # construction time (simulating a config where the router entry + # appears before its target deployments). The tier index must not be + # built eagerly — it's deferred until first use. + router = MagicMock() + router.model_list = [] # <- targets haven't been added yet + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku", "sonnet", "opus"]}, + ) + + # Now the targets come online. This mirrors the incremental add by + # `Router._create_deployment`. + router.model_list = _make_model_list( + [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": 2}, + {"model_name": "opus", "quality_tier": 3}, + ] + ) + + # First access triggers the index build and sees the full list. + assert qr._tier_to_models == { + 1: ["haiku"], + 2: ["sonnet"], + 3: ["opus"], + } + + +# ─── Router.set_model_list resets quality_routers (hot reload) ───────────── + + +class TestSetModelListResetsQualityRouters: + def test_set_model_list_clears_quality_routers_registry(self): + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + "model_info": {"litellm_routing_preferences": {"quality_tier": 1}}, + }, + { + "model_name": "my-qr", + "litellm_params": { + "model": "auto_router/quality_router", + "quality_router_default_model": "haiku", + "quality_router_config": {"available_models": ["haiku"]}, + }, + }, + ] + ) + + assert "my-qr" in router.quality_routers + + # Hot-reload with a new model_list that doesn't define the router. + router.set_model_list( + [ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + # Stale router from before must be cleared. + assert "my-qr" not in router.quality_routers + + +# ─── Pre-routing hook ─────────────────────────────────────────────────────── + + +class TestPreRoutingHook: + @pytest.mark.asyncio + async def test_simple_message_routes_to_tier_1(self, quality_router): + messages = [{"role": "user", "content": "hi"}] + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=messages, + ) + assert resp is not None + assert resp.model == "haiku" + + @pytest.mark.asyncio + async def test_reasoning_message_routes_to_tier_4(self, quality_router): + # Two reasoning markers triggers ComplexityTier.REASONING → quality 4. + messages = [ + { + "role": "user", + "content": ( + "Think step by step and reason through this problem. " + "Analyze this carefully and break down each component." + ), + } + ] + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=messages, + ) + assert resp is not None + assert resp.model == "opus-next" + + @pytest.mark.asyncio + async def test_empty_messages_returns_none(self, quality_router): + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=[], + ) + assert resp is None + + @pytest.mark.asyncio + async def test_only_system_message_routes_to_default(self, quality_router): + messages = [{"role": "system", "content": "You are a helpful assistant."}] + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=messages, + ) + assert resp is not None + assert resp.model == "haiku" # the configured default_model + + +# ─── Keyword override ────────────────────────────────────────────────────── + + +@pytest.fixture +def keyword_router(): + """ + Router where multiple deployments declare overlapping keywords so we can + exercise the (quality DESC, price ASC) tiebreak. + + - cheap-coder tier 2, keywords [code, python], cost 0.000001 + - smart-coder tier 3, keywords [code, python], cost 0.000010 + - law-bot tier 2, keywords [legal, contract], cost 0.000005 + - default-haiku tier 1, no keywords, cost 0.0000005 + """ + spec = [ + { + "model_name": "default-haiku", + "quality_tier": 1, + "keywords": [], + "input_cost_per_token": 0.0000005, + }, + { + "model_name": "cheap-coder", + "quality_tier": 2, + "keywords": ["code", "python"], + "input_cost_per_token": 0.000001, + }, + { + "model_name": "smart-coder", + "quality_tier": 3, + "keywords": ["code", "python"], + "input_cost_per_token": 0.000010, + }, + { + "model_name": "law-bot", + "quality_tier": 2, + "keywords": ["legal", "contract"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + return QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="default-haiku", + quality_router_config={ + "available_models": [ + "default-haiku", + "cheap-coder", + "smart-coder", + "law-bot", + ], + }, + ) + + +class TestKeywordOverride: + def test_no_keyword_in_message_returns_none(self, keyword_router): + assert keyword_router._keyword_override("hello there") is None + + def test_single_match_returns_that_model(self, keyword_router): + # Only law-bot declares "legal". + assert keyword_router._keyword_override("review this legal doc") == ( + "law-bot", + "legal", + ) + + def test_case_insensitive_match(self, keyword_router): + assert keyword_router._keyword_override("LEGAL question") == ( + "law-bot", + "legal", + ) + + def test_overlap_picks_highest_quality_tier(self, keyword_router): + # Both cheap-coder (tier 2) and smart-coder (tier 3) declare "code". + # Quality wins over price → smart-coder. + assert keyword_router._keyword_override("write some code for me") == ( + "smart-coder", + "code", + ) + + def test_same_tier_picks_cheapest(self): + # Two models at the same tier, both matching "data" — cheapest wins. + spec = [ + { + "model_name": "expensive", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000050, + }, + { + "model_name": "cheap", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="cheap", + quality_router_config={"available_models": ["expensive", "cheap"]}, + ) + match = qr._keyword_override("show me the data") + assert match == ("cheap", "data") + + def test_unpriced_loses_to_priced_at_same_tier(self): + # Same quality tier, one has cost, one doesn't → priced wins. + spec = [ + { + "model_name": "no-price", + "quality_tier": 2, + "keywords": ["data"], + # input_cost_per_token deliberately omitted + }, + { + "model_name": "with-price", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="no-price", + quality_router_config={"available_models": ["no-price", "with-price"]}, + ) + match = qr._keyword_override("show me the data") + assert match == ("with-price", "data") + + @pytest.mark.asyncio + async def test_hook_short_circuits_complexity_on_keyword_match( + self, keyword_router + ): + # A reasoning-style prompt would normally route to a high-quality model + # via the complexity flow — but the keyword "code" should short-circuit + # to smart-coder (highest tier among "code" models). + messages = [ + { + "role": "user", + "content": ( + "Think step by step and reason through this code problem. " + "Analyze this carefully and break down each component." + ), + } + ] + request_kwargs: Dict[str, Any] = {} + resp = await keyword_router.async_pre_routing_hook( + model="qr", + request_kwargs=request_kwargs, + messages=messages, + ) + assert resp is not None + assert resp.model == "smart-coder" + + decision = request_kwargs["metadata"]["quality_router_decision"] + assert decision["routed_via"] == "keyword" + assert decision["matched_keyword"] == "code" + assert decision["complexity_tier"] is None # short-circuited + + def test_quality_wins_over_explicit_order(self): + # Quality always beats order. A tier-3 model with no `order` wins over + # a tier-2 model with `order=1`. + spec = [ + { + "model_name": "ordered-tier2", + "quality_tier": 2, + "keywords": ["code"], + "order": 1, + "input_cost_per_token": 0.000010, + }, + { + "model_name": "implicit-tier3", + "quality_tier": 3, + "keywords": ["code"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="ordered-tier2", + quality_router_config={ + "available_models": ["ordered-tier2", "implicit-tier3"] + }, + ) + match = qr._keyword_override("write some code") + assert match == ("implicit-tier3", "code") + + def test_order_breaks_tie_within_same_quality_tier(self): + # Two tier-3 models, both match "code". Lower `order` wins. + spec = [ + { + "model_name": "preferred", + "quality_tier": 3, + "keywords": ["code"], + "order": 1, + "input_cost_per_token": 0.000050, # more expensive + }, + { + "model_name": "default-tier3", + "quality_tier": 3, + "keywords": ["code"], + "input_cost_per_token": 0.000005, # cheaper + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="default-tier3", + quality_router_config={"available_models": ["preferred", "default-tier3"]}, + ) + match = qr._keyword_override("write some code") + assert match == ("preferred", "code") + + def test_explicit_order_overrides_price(self): + # Same tier, but the more expensive one has a lower `order` and wins. + spec = [ + { + "model_name": "expensive-but-preferred", + "quality_tier": 2, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000050, + }, + { + "model_name": "cheap-default", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="cheap-default", + quality_router_config={ + "available_models": ["expensive-but-preferred", "cheap-default"] + }, + ) + match = qr._keyword_override("show me the data") + assert match == ("expensive-but-preferred", "data") + + def test_lower_order_wins_between_two_explicitly_ordered(self): + spec = [ + { + "model_name": "second", + "quality_tier": 2, + "keywords": ["data"], + "order": 5, + }, + { + "model_name": "first", + "quality_tier": 2, + "keywords": ["data"], + "order": 1, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="first", + quality_router_config={"available_models": ["first", "second"]}, + ) + match = qr._keyword_override("show me the data") + assert match == ("first", "data") + + def test_same_order_falls_through_to_quality_then_price(self): + # All three models share order=1 → tiebreak falls through to + # (quality DESC, cost ASC). + spec = [ + { + "model_name": "low-tier", + "quality_tier": 1, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000001, + }, + { + "model_name": "high-tier-cheap", + "quality_tier": 3, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000005, + }, + { + "model_name": "high-tier-expensive", + "quality_tier": 3, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000050, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="low-tier", + quality_router_config={ + "available_models": [ + "low-tier", + "high-tier-cheap", + "high-tier-expensive", + ] + }, + ) + match = qr._keyword_override("show me the data") + assert match == ("high-tier-cheap", "data") + + def test_order_is_used_in_tier_resolution_too(self): + # Two models at the same tier. Explicit `order=1` on the second one + # should make _resolve_model_for_quality_tier(2) pick it. + spec = [ + { + "model_name": "default-pick", + "quality_tier": 2, + }, + { + "model_name": "preferred-pick", + "quality_tier": 2, + "order": 1, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="default-pick", + quality_router_config={ + "available_models": ["default-pick", "preferred-pick"] + }, + ) + assert qr._resolve_model_for_quality_tier(2) == "preferred-pick" + + @pytest.mark.asyncio + async def test_hook_falls_back_to_complexity_when_no_keyword(self, keyword_router): + # No declared keyword in the message → complexity-based routing. + # "hi" is SIMPLE → quality 1 → default-haiku (the only tier-1 model). + messages = [{"role": "user", "content": "hi"}] + request_kwargs: Dict[str, Any] = {} + resp = await keyword_router.async_pre_routing_hook( + model="qr", + request_kwargs=request_kwargs, + messages=messages, + ) + assert resp is not None + assert resp.model == "default-haiku" + + decision = request_kwargs["metadata"]["quality_router_decision"] + assert decision["routed_via"] == "quality_tier" + assert decision["matched_keyword"] is None + assert decision["complexity_tier"] == "SIMPLE" + + +# ─── Routing-decision metadata (powers x-litellm-quality-router-* headers) ── + + +class TestDecisionMetadata: + @pytest.mark.asyncio + async def test_hook_stashes_decision_in_request_kwargs_metadata( + self, quality_router + ): + # Reasoning prompt → REASONING → quality tier 4 → opus-next. + messages = [ + { + "role": "user", + "content": ( + "Think step by step and reason through this problem. " + "Analyze this carefully and break down each component." + ), + } + ] + request_kwargs: Dict[str, Any] = {} + + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs=request_kwargs, + messages=messages, + ) + assert resp is not None and resp.model == "opus-next" + + decision = request_kwargs["metadata"]["quality_router_decision"] + assert decision["routed_model"] == "opus-next" + assert decision["quality_tier"] == 4 + assert decision["complexity_tier"] == "REASONING" + assert decision["router_model_name"] == "quality-router-test" + assert decision["routed_via"] == "quality_tier" + assert decision["matched_keyword"] is None + + @pytest.mark.asyncio + async def test_decision_metadata_preserves_existing_metadata(self, quality_router): + request_kwargs: Dict[str, Any] = { + "metadata": {"trace_id": "abc-123", "user_id": "u-1"} + } + + await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + # Existing metadata keys are intact and the decision is added alongside. + assert request_kwargs["metadata"]["trace_id"] == "abc-123" + assert request_kwargs["metadata"]["user_id"] == "u-1" + assert "quality_router_decision" in request_kwargs["metadata"] + + +# ─── Router.set_response_headers lifts decision into x-litellm-quality-* ──── + + +class TestSetResponseHeadersLiftsDecision: + """ + Verify the Router.set_response_headers helper turns a stashed quality-router + decision into x-litellm-quality-router-* headers on the response. + """ + + @pytest.mark.asyncio + async def test_lifts_decision_into_additional_headers(self): + from pydantic import BaseModel + + from litellm.router import Router + + class FakeResponse(BaseModel): + model_config = {"arbitrary_types_allowed": True} + _hidden_params: Dict[str, Any] = {} + + # Build a real Router with a tiny model_list — enough to satisfy + # set_response_headers without needing the rest of the router stack. + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + response = FakeResponse() + response._hidden_params = {} + + request_kwargs = { + "metadata": { + "quality_router_decision": { + "router_model_name": "qr", + "routed_model": "smart-coder", + "routed_via": "keyword", + "matched_keyword": "code", + "quality_tier": 3, + "complexity_tier": None, + } + } + } + + await router.set_response_headers( + response=response, + model_group="qr", + request_kwargs=request_kwargs, + ) + + headers = response._hidden_params["additional_headers"] + assert headers["x-litellm-quality-router-model"] == "smart-coder" + assert headers["x-litellm-quality-router-tier"] == "3" + assert headers["x-litellm-quality-router-via"] == "keyword" + assert headers["x-litellm-quality-router-keyword"] == "code" + # Keyword route short-circuits classification → no complexity header. + assert "x-litellm-quality-router-complexity" not in headers + # Existing x-litellm-model-group behavior is unchanged. + assert headers["x-litellm-model-group"] == "qr" + + @pytest.mark.asyncio + async def test_quality_tier_route_emits_complexity_not_keyword(self): + from pydantic import BaseModel + + from litellm.router import Router + + class FakeResponse(BaseModel): + model_config = {"arbitrary_types_allowed": True} + _hidden_params: Dict[str, Any] = {} + + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + response = FakeResponse() + response._hidden_params = {} + + request_kwargs = { + "metadata": { + "quality_router_decision": { + "router_model_name": "qr", + "routed_model": "haiku", + "routed_via": "quality_tier", + "matched_keyword": None, + "quality_tier": 1, + "complexity_tier": "SIMPLE", + } + } + } + + await router.set_response_headers( + response=response, + model_group="qr", + request_kwargs=request_kwargs, + ) + + headers = response._hidden_params["additional_headers"] + assert headers["x-litellm-quality-router-via"] == "quality_tier" + assert headers["x-litellm-quality-router-complexity"] == "SIMPLE" + # Quality-tier route → no keyword header. + assert "x-litellm-quality-router-keyword" not in headers + + @pytest.mark.asyncio + async def test_no_decision_leaves_quality_router_headers_unset(self): + from pydantic import BaseModel + + from litellm.router import Router + + class FakeResponse(BaseModel): + model_config = {"arbitrary_types_allowed": True} + _hidden_params: Dict[str, Any] = {} + + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + response = FakeResponse() + response._hidden_params = {} + + await router.set_response_headers( + response=response, + model_group="haiku", + request_kwargs={}, # no quality_router_decision + ) + + headers = response._hidden_params["additional_headers"] + assert "x-litellm-quality-router-model" not in headers + assert "x-litellm-quality-router-tier" not in headers + + +class TestRouterQualityDeploymentMethods: + """Tests for Router._is_quality_router_deployment and Router.init_quality_router_deployment.""" + + def test_is_quality_router_deployment_true(self): + """_is_quality_router_deployment returns True for quality router models.""" + from litellm.router import Router + from litellm.types.router import LiteLLM_Params + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + params = LiteLLM_Params(model="auto_router/quality_router/my-router") + assert router._is_quality_router_deployment(params) is True + + def test_is_quality_router_deployment_false(self): + """_is_quality_router_deployment returns False for regular models.""" + from litellm.router import Router + from litellm.types.router import LiteLLM_Params + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + params = LiteLLM_Params(model="openai/gpt-4o-mini") + assert router._is_quality_router_deployment(params) is False + + def test_init_quality_router_deployment(self): + """init_quality_router_deployment registers a QualityRouter.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + deployment = Deployment( + model_name="auto_router/quality_router/test-router", + litellm_params=LiteLLM_Params( + model="auto_router/quality_router/test-router", + quality_router_default_model="gpt-4o-mini", + ), + model_info={"id": "test-id"}, + ) + router.init_quality_router_deployment(deployment) + assert "auto_router/quality_router/test-router" in router.quality_routers