From eac5dc10f338e7faf751dd198d10e0bd7a2c024b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:11:36 -0700 Subject: [PATCH 1/6] fix(guardrails): apply PUT /guardrails/{id} to the serving worker immediately and reject invalid configs with 422 --- .../proxy/guardrails/guardrail_endpoints.py | 20 +++- .../proxy/guardrails/guardrail_registry.py | 37 ++---- .../guardrails/test_guardrail_endpoints.py | 53 +++++++-- .../guardrails/test_guardrail_registry.py | 108 ++++++++++++++---- 4 files changed, 158 insertions(+), 60 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b04828f0f2..3d2ed641a30 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -514,12 +514,26 @@ async def update_guardrail( guardrail_name: Final = result.get("guardrail_name", "Unknown") try: - IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail( - guardrail_id=guardrail_id, guardrail=cast(Guardrail, result) - ) + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result)) verbose_proxy_logger.info( "Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id ) + except (ValueError, TypeError) as update_error: + # The new config is invalid (a raising guardrail __init__): + # reinitialize_guardrail already restored the previous live instance, but + # update_guardrail_in_db above already persisted the rejected config to + # the DB. Roll that back too, so the DB and the live guardrail never + # disagree about what's actually enforcing, and surface the rejection to + # the caller instead of a misleading 200. + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=existing_guardrail, + prisma_client=prisma_client, + ) + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {update_error}", + ) from update_error except Exception as update_error: verbose_proxy_logger.warning( "Immediate sync: Failed to update '%s' (ID: %s) in memory: %s", diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index dc13c09dd38..60873a1eeb1 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -6,7 +6,7 @@ import os from collections.abc import Callable, Iterator, Mapping from datetime import datetime, timezone from itertools import chain, count -from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol from pydantic import ValidationError @@ -615,26 +615,6 @@ class InMemoryGuardrailHandler: return _guardrail_callback - def update_in_memory_guardrail( - self, - guardrail_id: str, - guardrail: Guardrail, - source: Literal["db", "config"] = "db", - ) -> None: - """ - Update a guardrail in memory - - - updates the guardrail in memory - - updates the guardrail params in litellm.callback_manager - """ - self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail - self._sources[guardrail_id] = source - - custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id) - if custom_guardrail_callback: - updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) - custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) - def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ Delete a guardrail in memory and remove from litellm callbacks. @@ -789,11 +769,12 @@ class InMemoryGuardrailHandler: Removes old callback from litellm.callbacks and creates fresh instance. If the new config fails to initialize (e.g. an invalid on_flagged - combination), the previous instance is restored rather than left - deleted: initialize_guardrail's own ValueError/TypeError propagate - uncaught, so a caller reaching this point after already deleting the - old instance would otherwise leave the guardrail providing no - protection at all, not merely "still enforcing the old config." + combination or an invalid regex), the previous instance is restored + rather than left deleted, and the failure is re-raised as ValueError so + every init failure reaches callers as one exception type: a caller + reaching this point after already deleting the old instance would + otherwise leave the guardrail providing no protection at all, not + merely "still enforcing the old config." """ guardrail_id: Final = guardrail.get("guardrail_id") if not guardrail_id: @@ -812,7 +793,7 @@ class InMemoryGuardrailHandler: # that was enforcing must never fail open because an update was bad. try: return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source) - except Exception: + except Exception as init_error: if previous_guardrail is not None: verbose_proxy_logger.exception( "Reinitializing guardrail %s with updated params failed; restoring the previous configuration", @@ -824,7 +805,7 @@ class InMemoryGuardrailHandler: ) except Exception: # noqa: BLE001 # the original failure must propagate even if the restore breaks verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id) - raise + raise ValueError(f"Guardrail initialization failed: {init_error}") from init_error def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None: """ diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 9b2117b7647..320e51203f6 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -104,7 +104,7 @@ def mock_in_memory_handler(mocker): mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL mock_handler.get_source.return_value = "config" mock_handler.initialize_guardrail = mocker.Mock() - mock_handler.update_in_memory_guardrail = mocker.Mock() + mock_handler.sync_guardrail_from_db = mocker.Mock() mock_handler.delete_in_memory_guardrail = mocker.Mock() mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[]) return mock_handler @@ -1045,13 +1045,15 @@ async def test_create_guardrail_endpoint( "scenario,expected_result,expected_exception", [ ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), + ("success_sync_fails_unexpected_error", "test-db-guardrail", None), + ("sync_fails_invalid_config", None, HTTPException), ("database_failure", None, HTTPException), ("no_prisma_client", None, HTTPException), ], ids=[ "success_with_immediate_sync", - "success_but_sync_fails", + "success_but_sync_fails_with_unexpected_error", + "sync_rejects_invalid_config", "database_error", "missing_prisma_client", ], @@ -1071,6 +1073,7 @@ async def test_update_guardrail_endpoint( mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", @@ -1081,10 +1084,13 @@ async def test_update_guardrail_endpoint( mock_in_memory_handler, ) - elif scenario == "success_sync_fails": + elif scenario == "success_sync_fails_unexpected_error": + # A non-ValueError/TypeError failure is not a config-rejection signal, + # so it keeps the pre-existing swallow-and-warn behavior rather than + # rolling back the DB write. mock_prisma_client = mocker.Mock() - mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception( - "Sync failed" + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=Exception("Sync failed") ) mock_logger = mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" @@ -1100,6 +1106,25 @@ async def test_update_guardrail_endpoint( mock_in_memory_handler, ) + elif scenario == "sync_fails_invalid_config": + # Regression for the PUT half of the fix: a TypeError from the sync (the + # deleted update_in_memory_guardrail raised exactly this on every PUT) + # must roll back the DB write and surface a 422, not persist the + # rejected config with a 200. + mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=TypeError("vars() argument must have __dict__ attribute") + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern + mocker.patch( # test-quality-ok: reused pattern + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( # test-quality-ok: reused pattern + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( @@ -1128,6 +1153,16 @@ async def test_update_guardrail_endpoint( assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) + elif scenario == "sync_fails_invalid_config": + assert exc_info.value.status_code == 422 + assert "update rejected" in str(exc_info.value.detail) + # Rolled back: update_guardrail_in_db is called once for the + # rejected write and once more to restore the previous config. + assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2 + assert ( + mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"] + == MOCK_DB_GUARDRAIL + ) else: result = await update_guardrail( @@ -1143,11 +1178,11 @@ async def test_update_guardrail_endpoint( prisma_client=mocker.ANY, ) - mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with( - guardrail_id="test-guardrail-id", guardrail=mocker.ANY + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( + guardrail=mocker.ANY ) - if scenario == "success_sync_fails": + if scenario == "success_sync_fails_unexpected_error": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 2c0735970d3..beaffa73100 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -154,29 +154,51 @@ def test_duplicate_config_guardrail_names_get_distinct_stable_ids(): registry_module.guardrail_initializer_registry.pop("dup_name_test", None) -def test_update_in_memory_guardrail(): +def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance(): + """ + Regression for PUT /guardrails/{id}: the DB row arrives with litellm_params as + a plain jsonb dict, and the deleted update_in_memory_guardrail cast it to + LitellmParams without constructing one, so vars() raised and the running proxy + kept enforcing the stale config forever. The PUT endpoint now routes through + sync_guardrail_from_db, which must rebuild the live instance from the dict: + new blocked words compiled in, old ones gone, and the event hook re-derived + from mode (the base-class setattr path wrote self.mode while dispatch reads + self.event_hook, so only a full re-init applies a mode change). + """ + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + handler = InMemoryGuardrailHandler() - handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail( - guardrail_name="test-guardrail", - default_on=False, - event_hook=GuardrailEventHooks.pre_call, - ) + gid = "66666666-6666-6666-6666-666666666666" - handler.update_in_memory_guardrail( - "123", - Guardrail( - guardrail_name="test-guardrail", - litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True), - ), - ) - - assert ( - handler.guardrail_id_to_custom_guardrail["123"].should_run_guardrail( - data={}, event_type=GuardrailEventHooks.pre_call + def db_guardrail(word: str, mode: str) -> Guardrail: + return Guardrail( + guardrail_id=gid, + guardrail_name="cf-put-sync", + litellm_params={ + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": True, + "blocked_words": [{"keyword": word, "action": "BLOCK"}], + }, ) - is True - ) - assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.sync_guardrail_from_db(db_guardrail("foobarblock", "pre_call")) + handler.sync_guardrail_from_db(db_guardrail("quxnewblock", "during_call")) + + instance = handler.guardrail_id_to_custom_guardrail[gid] + assert isinstance(instance, ContentFilterGuardrail) + assert instance._check_blocked_words("hello QUXNEWBLOCK") is not None + assert instance._check_blocked_words("hello FOOBARBLOCK") is None + assert instance.event_hook == GuardrailEventHooks.during_call + assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: @@ -774,3 +796,49 @@ def test_reinitialize_guardrail_restores_previous_on_failure(): assert restored.guardrail_name == "restore-me" finally: registry_module.guardrail_initializer_registry.pop("restore_test", None) + + +def test_reinitialize_guardrail_raises_value_error_for_non_value_error_init_failures(): + """Regression for the LIT-6479 fix's 422 path: a constructor failure that is not + already a ValueError/TypeError (re.error from an invalid regex has neither in its + MRO) must still surface as ValueError, so the PUT/PATCH endpoints' rollback+422 + catch is exhaustive instead of warn-and-200 persisting a broken config.""" + import re + + from litellm.proxy.guardrails import guardrail_registry as registry_module + + def _initializer(litellm_params, guardrail): + if litellm_params.api_key == "bad-regex": + re.compile("([") + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + registry_module.guardrail_initializer_registry["regex_test"] = _initializer + try: + handler = InMemoryGuardrailHandler() + created = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "regex-me", + "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "ok"}, + }, + ) + guardrail_id = created["guardrail_id"] + + with pytest.raises(ValueError, match="Guardrail initialization failed") as excinfo: + handler.reinitialize_guardrail( + guardrail={ + "guardrail_id": guardrail_id, + "guardrail_name": "regex-me", + "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "bad-regex"}, + }, + ) + + assert isinstance(excinfo.value.__cause__, re.error) + assert guardrail_id in handler.IN_MEMORY_GUARDRAILS + restored = handler.guardrail_id_to_custom_guardrail[guardrail_id] + assert restored is not None and restored.guardrail_name == "regex-me" + finally: + registry_module.guardrail_initializer_registry.pop("regex_test", None) From d33fe95d192a819eef63848468c3630903a135a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:08:45 -0700 Subject: [PATCH 2/6] test(responses): expect the 404 OpenAI now returns for an unknown model --- .../test_e2e_openai_responses_api.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 1a7fb1f3e41..abae26e02cd 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,5 +1,5 @@ import httpx -from openai import OpenAI, BadRequestError, APIStatusError +from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError import pytest @@ -105,10 +105,9 @@ def test_streaming_response(): assert len(collected_chunks) > 0 -def test_bad_request_error(): +def test_model_not_found_error(): client = get_test_client() - with pytest.raises(BadRequestError): - # Trigger error with invalid model name + with pytest.raises(NotFoundError): client.responses.create(model="non-existent-model", input="This should fail") From 87618b7da8e755cb007b3bf2470dd05b93154f6d Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 2 Sep 2026 18:12:03 -0700 Subject: [PATCH 3/6] doc (#39456) --- tests/rust-python-harness/AGENTS.md | 44 +++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 tests/rust-python-harness/AGENTS.md diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md new file mode 100644 index 00000000000..e9b17027ddc --- /dev/null +++ b/tests/rust-python-harness/AGENTS.md @@ -0,0 +1,44 @@ +# Expected Structure + +```text +tests/rust-python-harness/ +├── __main__.py +│ +├── strategies/ +│ ├── e2e_parity/ +│ │ ├── runner.py +│ │ ├── sdk/ +│ │ │ ├── ocr/ +│ │ │ ├── messages/ +│ │ │ ├── chat_completions/ +│ │ │ └── responses/ +│ │ └── gateway/ +│ │ +│ ├── trace_parity/ +│ │ ├── runner.py +│ │ ├── sdk/ +│ │ └── gateway/ +│ │ +│ └── unit_tests/ +│ ├── runner.py +│ ├── mapping_validator.py +│ ├── python_runner.py +│ └── rust_runner.py +│ +└── shared/ + ├── parity/ + ├── tracing/ + └── reporting/ +``` + +- Run locally only; no CI integration +- `__main__.py` selects strategies and combines their reports; each strategy also runs independently +- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses +- `trace_parity/` compares mapped operations, call counts, and required execution ordering +- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders +- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs +- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts +- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results +- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation +- `shared/` contains reusable parity, tracing, and reporting machinery +- Keep fixtures with their owning API and existing Python tests in their current locations From 4f7b20ec102bc4f01f32156bc9d0d8325af3a081 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:25:28 -0700 Subject: [PATCH 4/6] fix(guardrails): skip streaming guardrail rounds that re-scan cleared output (#39386) * fix(guardrails): skip streaming guardrail rounds that re-scan cleared output Streaming guardrails scanned the finished answer twice at end of stream whenever the chunk count landed on a multiple of the sampling rate, ran sampled rounds whose payload was identical to the previous one, and on /v1/messages could scan an empty text before the first content chunk. Every redundant round is a paid guardrail provider call. Each endpoint handler now exposes a scan key describing what a round would hand to apply_guardrail (the text so far, plus tool calls once the stream has ended), and the unified streaming hook skips a sampled or end-of-stream round whose key equals the last scanned one or carries nothing to scan yet. Rounds that carry tool calls are never skipped. * test(guardrails): expect one end-of-stream scan when the terminal chunk is sampled Update sampled cadence expectations and use tuple-backed scan state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../a2a/chat/guardrail_translation/handler.py | 12 +- .../chat/guardrail_translation/handler.py | 25 +- .../guardrail_translation/base_translation.py | 19 ++ .../base_llm/guardrail_translation/utils.py | 12 + .../chat/guardrail_translation/handler.py | 37 ++- .../guardrail_translation/handler.py | 72 ++++- .../unified_guardrail/unified_guardrail.py | 40 ++- .../test_a2a_guardrail_handler.py | 35 +++ .../test_anthropic_guardrail_handler.py | 54 ++++ .../test_openai_guardrail_handler.py | 72 +++++ ...test_openai_responses_guardrail_handler.py | 90 ++++++ .../test_openai_moderation_streaming.py | 19 +- .../test_generic_guardrail_api.py | 18 +- .../test_unified_guardrail.py | 268 ++++++++++++++++++ 14 files changed, 733 insertions(+), 40 deletions(-) create mode 100644 tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index f1c7451796d..5c30ff4747a 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation): return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + _, valid_parsed = self._parse_streaming_responses(responses_so_far) + combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed) + return StreamingScanKey(texts=(combined_text,)) + def _parse_streaming_responses( self, - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c7d12e5cf3a..c23797f72af 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( anthropic_tool_name, anthropic_tool_names, @@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @classmethod + def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(block) + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..b28daf73bc4 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -35,6 +35,22 @@ class StreamTransformSink: holdback_per_choice: dict[int, int] = field(default_factory=dict) +@dataclass(frozen=True, slots=True) +class StreamingScanKey: + """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys + compare equal when the round would scan the same content again; ``stream_ended`` + stays out of the comparison and only says whether the handler is on its + end-of-stream path, where an empty payload is still scanned today.""" + + texts: tuple[str, ...] + tool_calls: tuple[str, ...] = () + stream_ended: bool = field(default=False, compare=False) + + @property + def has_nothing_to_scan(self) -> bool: + return not self.stream_ended and not any(self.texts) and not self.tool_calls + + class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( @@ -151,6 +167,9 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + return None + def build_block_sse_chunks( self, exc: "ModifyResponseException", diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 9b6f9c47105..8dee262001d 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,6 +4,8 @@ import json from collections.abc import Callable, Iterator, Sequence from typing import Any, Final, TypeVar +from pydantic import BaseModel + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None: return getattr(item, field, None) +def stream_item_fingerprint(item: object) -> str: + plain: Final = item.model_dump() if isinstance(item, BaseModel) else item + return json.dumps(plain, sort_keys=True, default=str) + + +def stream_item_items(item: object, field: str) -> tuple[object, ...]: + value: Final = stream_item_field(item, field) + return tuple(value) if isinstance(value, (list, tuple)) else () + + def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: """ ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 96a5ed663fc..d41c8557d72 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,6 +26,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( role_out_of_guardrail_scope, scoped_structured_message_indices, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here (see ``_process_streaming_transform`` for the incremental_diff path).""" - # check if the stream has ended - has_stream_ended = False - for chunk in responses_so_far: - if chunk.choices and chunk.choices[0].finish_reason is not None: - has_stream_ended = True - break + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: # convert to model response @@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) } + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + return StreamingScanKey( + texts=tuple(self._combine_streaming_texts(chunks).values()), + tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @staticmethod + def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(tool_call) + for chunk in responses_so_far + for choice in _stream_chunk_choices(chunk) + for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + ) + + @staticmethod + def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool: + first_choices: Final = tuple( + choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices + ) + return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices) + def _combine_streaming_texts( - self, responses_so_far: list["ModelResponseStream"] + self, responses_so_far: Sequence["ModelResponseStream"] ) -> dict[tuple[int, int | None], str]: """ Combine all streaming chunks into complete text per choice. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 5a5970fb867..a0db7aadb9e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation): ) return responses_so_far - def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + terminal_types: Final = frozenset( + ( + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + ) + ) + return stream_item_field(responses_so_far[-1], "type") in terminal_types + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + if not responses_so_far or not hasattr(responses_so_far[-1], "get"): + return None + last_event: Final = responses_so_far[-1] + last_event_type: Final = stream_item_field(last_event, "type") + if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: + return None + if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + return self._completed_response_scan_key(stream_item_field(last_event, "response")) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + stream_ended=self._check_streaming_has_ended(responses_so_far), + ) + + @staticmethod + def _completed_response_scan_key(response: object) -> StreamingScanKey: + output_items: Final = stream_item_items(response, "output") + message_items: Final = tuple( + item for item in output_items if stream_item_field(item, "type") != "function_call" + ) + return StreamingScanKey( + texts=tuple( + text + for item in message_items + for part in stream_item_items(item, "content") + if isinstance(text := stream_item_field(part, "text"), str) and text + ), + tool_calls=tuple( + stream_item_fingerprint(item) + for item in output_items + if stream_item_field(item, "type") == "function_call" + ), + stream_ended=True, + ) def build_stream_error_items( self, @@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation): ), ) - def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Get the string so far from the responses so far. @@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation): """ keyed_events: Final = tuple( ( - (event.get("item_id"), event.get("output_index"), event.get("content_index")), - event.get("text"), - event.get("delta"), + ( + stream_item_field(event, "item_id"), + stream_item_field(event, "output_index"), + stream_item_field(event, "content_index"), + ), + stream_item_field(event, "text"), + stream_item_field(event, "delta"), ) for event in responses_so_far - if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str) ) def part_text(part_key: tuple[object, object, object]) -> str: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 46b00829b74..c6b8df1b493 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, ) # Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error @@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol): @property def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + @property + def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... + @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... @@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]: return choices +def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool: + if scan_key is None: + return False + return scan_key == last_scan_key or scan_key.has_nothing_to_scan + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" @@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False + last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round async for item in response: chunk_counter += 1 @@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", + chunk_counter, + guardrail_to_apply.guardrail_name, + ) + chunks_yielded = True + responses_yielded.append(item) + yield item + continue + verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, @@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger): # string, permanently losing this chunk's content. original_item = copy.deepcopy(item) - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, @@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger): ): yield error_item return + if scan_key is not None: + last_scan_key = scan_key chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger): # preserve the list, not clone every chunk (deepcopy would double # peak memory for large responses). buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(end_scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", + guardrail_to_apply.guardrail_name, + ) + for buffered_item in buffered_items or (): + yield buffered_item + for pending_item in pending_end_of_stream_items: + responses_yielded.append(pending_item) + yield pending_item + return try: await endpoint_translation.process_output_streaming_response( diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py new file mode 100644 index 00000000000..dd7e8fadcd8 --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -0,0 +1,35 @@ +"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" + +import json + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + +def _text_event(text: str) -> str: + return json.dumps( + { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]}, + } + ) + + +def _status_event() -> str: + return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}}) + + +class TestA2AGuardrailHandlerStreamingScanKey: + def test_key_joins_the_text_of_every_message_event(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")]) + assert key == StreamingScanKey(texts=("hello world",)) + + def test_events_without_text_leave_the_key_unchanged(self): + handler = A2AGuardrailHandler() + events = [_text_event("hello")] + assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events) + + def test_unparseable_items_are_ignored(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 0fe7730e91e..3044a321aa6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) @@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults: } later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)] assert {"type": "text", "text": "Now fetch the page."} in later_blocks + + +class TestAnthropicMessagesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _sse(event_type, data): + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + def _text_delta(self, text): + return self._sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + def test_key_is_empty_before_any_text_arrives(self): + head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}}) + key = AnthropicMessagesHandler().get_streaming_scan_key([head]) + assert key == StreamingScanKey(texts=("",)) + + def test_key_accumulates_text_deltas(self): + key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")]) + assert key.texts == ("hello world",) + assert key.stream_ended is False + + def _stop(self, stop_reason): + return self._sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}}, + ) + + def test_stop_without_tool_use_scans_the_same_payload(self): + handler = AnthropicMessagesHandler() + open_key = handler.get_streaming_scan_key([self._text_delta("hi")]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self): + handler = AnthropicMessagesHandler() + tool_use = self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}, + }, + ) + open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) + assert open_key == StreamingScanKey(texts=("hi",)) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key 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 7dd6065063a..cebab2512d0 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 @@ -12,6 +12,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) @@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded: ) ] assert handler._check_streaming_has_ended(chunks) is True + + +class TestStreamingScanKey: + """get_streaming_scan_key identifies what a sampled round would scan so the + unified hook can skip rounds that would re-scan already-cleared text""" + + @staticmethod + def _chunk(content, finish_reason=None, index=0): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + def test_key_carries_accumulated_text_and_open_stream(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_chunks_without_text_leave_the_key_unchanged(self): + handler = OpenAIChatCompletionsHandler() + before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)]) + assert after == before + + def test_finish_chunk_without_tool_calls_scans_the_same_payload(self): + handler = OpenAIChatCompletionsHandler() + open_key = handler.get_streaming_scan_key([self._chunk("hi")]) + ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")]) + assert open_key.stream_ended is False + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self): + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + handler = OpenAIChatCompletionsHandler() + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + tool_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] + ) + assert open_key == StreamingScanKey(texts=("hi",)) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): + handler = OpenAIChatCompletionsHandler() + first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] + key_at_first_finish = handler.get_streaming_scan_key(first_done) + key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)]) + assert key_at_first_finish.stream_ended is True + assert key_after_more_text.stream_ended is True + assert key_after_more_text != key_at_first_finish + + def test_non_stream_items_are_ignored(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"]) + assert key.texts == ("hi",) 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 d071ef78c2d..295121167d6 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 @@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks: dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] assert len(dones) == 1 assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." + + +class TestOpenAIResponsesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + def test_no_events_yields_no_key(self): + assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None + + def test_key_accumulates_deltas_while_the_stream_is_open(self): + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_typed_delta_events_accumulate_like_dicts(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + events = [ + OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=i, + ) + for i, text in enumerate(("hel", "lo")) + ] + key = OpenAIResponsesHandler().get_streaming_scan_key(events) + assert key.texts == ("hello",) + assert key.stream_ended is False + + def test_events_without_text_leave_the_key_unchanged(self): + handler = OpenAIResponsesHandler() + events = [self._delta(0, "hi")] + quiet = events + [{"type": "response.in_progress", "sequence_number": 1}] + assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events) + + @staticmethod + def _completed(sequence_number, output): + return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}} + + def test_completed_event_keys_on_the_final_output_text(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_completed_event_with_a_function_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])]) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_completed_event_reads_every_output_text_part(self): + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + item = GenericResponseOutputItem( + type="message", + id="msg_1", + status="completed", + role="assistant", + content=[ + OutputText(type="output_text", text="one", annotations=[]), + OutputText(type="output_text", text="two", annotations=[]), + ], + ) + key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])]) + assert key.texts == ("one", "two") + + def test_output_item_done_round_is_never_deduped(self): + done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} + assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 476d443d8d8..cb6772977ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None): @pytest.mark.asyncio async def test_openai_moderation_streaming_default_uses_sampled_cadence(): """Default config samples every 5th streamed chunk and runs a final aggregate - pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream - calls, plus 1 final = 3 total. + pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream + round is skipped because chunk 10 already scanned the full text, for 2 total calls """ import litellm @@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence(): ): pass - assert patched_make_request.await_count == 3, ( - f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), " + assert patched_make_request.await_count == 2, ( + f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {patched_make_request.await_count}" ) @@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder @pytest.mark.asyncio async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled(): """With streaming_end_of_stream_only=False and streaming_sampling_rate=2, - moderation runs every 2nd chunk during the stream, plus once more at end. + moderation runs every 2nd chunk during the stream. The terminal chunk scan covers + the final aggregate, for 3 total calls """ import litellm @@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab ): pass - # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls), - # plus the final aggregate pass after the stream ends (1 call) = 4 total. - assert patched_make_request.await_count == 4, ( - f"Expected 4 moderation calls (3 sampled + 1 final aggregate), " + assert patched_make_request.await_count == 3, ( + f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {patched_make_request.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 523ec1a37b4..83cc9ae8bb9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_default_uses_sampled_cadence(self): - """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3.""" + """Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10 + cover the full text, so the end-of-stream round is skipped and there are 2 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 3, ( - f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), " + assert mock_post.await_count == 2, ( + f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {mock_post.await_count}" ) for call in mock_post.await_args_list: @@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_sampling_rate_override(self): - """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls.""" + """sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so + the end-of-stream round is skipped and there are 3 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 4, ( - f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), " + assert mock_post.await_count == 3, ( + f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {mock_post.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8cad1c634a9..a28a2a71613 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket: assert recorded[0]["guardrail_name"] == "audit-recorder" assert recorded[0]["guardrail_status"] == "success" assert request_data["metadata"]["user_api_key_user_id"] == "user-1" + + +class _ScanCountingGuardrail(CustomGuardrail): + """Pass-through guardrail that records every response-side scan payload.""" + + def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False): + super().__init__(guardrail_name="scan-counter") + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = end_of_stream_only + self.streaming_buffer_until_moderated = buffer_until_moderated + self.guardrail_config = {} + self.scans: tuple[dict[str, object], ...] = () + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.scans = ( + *self.scans, + { + "texts": list(inputs.get("texts") or []), + "tool_calls": list(inputs.get("tool_calls") or []), + "model": inputs.get("model"), + }, + ) + return inputs + + +def _responses_delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + +def _responses_tail(sequence_number, text): + return [ + { + "type": "response.output_text.done", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": text, + }, + { + "type": "response.completed", + "sequence_number": sequence_number + 1, + "response": { + "model": "gpt-5.6", + "output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}], + }, + }, + ] + + +class TestStreamingScanDedup: + """A sampled round whose scan payload matches the previous round (or carries + no text yet) is skipped, so a stream is never re-scanned for output the + guardrail already cleared. Regression for LIT-6692.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self, monkeypatch): + monkeypatch.setattr( + unified_module, + "endpoint_guardrail_translation_mappings", + load_guardrail_translation_mappings(), + ) + + @pytest.mark.asyncio + async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 3 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a"), + _stream_chunk("b"), + _stream_chunk("c"), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk("d", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self): + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + guardrail = _ScanCountingGuardrail(sampling_rate=3) + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + finish = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls") + ] + ) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a", index=0), + _stream_chunk("x", index=1), + _stream_chunk("b", finish_reason="stop", index=0), + _stream_chunk("y", index=1), + _stream_chunk("z", finish_reason="stop", index=1), + ] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(guardrail.scans) == 2 + assert any("yz" in text for text in guardrail.scans[-1]["texts"]) + + @pytest.mark.asyncio + async def test_responses_completed_event_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(8)] + full_text = "".join(f"t{i}" for i in range(8)) + chunks = deltas + _responses_tail(8, full_text) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 10 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]] + assert guardrail.scans[-1]["model"] == "gpt-5.6" + + @pytest.mark.asyncio + async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + chunks = deltas + _responses_tail(5, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]] + + @pytest.mark.asyncio + async def test_responses_completed_carrying_a_function_call_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + completed = { + "type": "response.completed", + "sequence_number": 5, + "response": { + "model": "gpt-5.6", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + }, + } + chunks = deltas + [completed] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 6 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_responses_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)] + chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 12 + assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},) + + @pytest.mark.asyncio + async def test_responses_tool_call_done_event_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + tool_call_done = { + "type": "response.output_item.done", + "sequence_number": 1, + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + } + chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 4 + assert len(guardrail.scans) == 2 + assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"] + assert guardrail.scans[1]["texts"] == ["hi"] + + @pytest.mark.asyncio + async def test_anthropic_skips_empty_round_and_terminal_duplicate(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_still_scans_exactly_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self): + guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] From 92edcb90dbb41175f526187826f70058e5a57f94 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:27:19 -0700 Subject: [PATCH 5/6] fix: keep litellm importable on Python 3.10 and guard 3.11-only typing imports in CI (#39448) * ci: guard against Python 3.10-incompatible typing imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: address Python 3.10 typing guard review Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ci): honor version-guard direction and scan litellm-proxy-extras in py310 typing check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-code-quality.yml | 33 ++++ .../websearch_interception/handler.py | 4 +- .../_experimental/mcp_server/tool_search.py | 4 +- litellm/proxy/agent_endpoints/endpoints.py | 4 +- .../proxy/common_utils/reset_budget_job.py | 4 +- .../batch_file_validation.py | 4 +- .../general_upload_validation.py | 4 +- .../types/llms/gemini_audio_transcription.py | 4 +- .../check_py310_typing_imports.py | 150 ++++++++++++++++++ .../test_check_py310_typing_imports.py | 86 ++++++++++ 10 files changed, 286 insertions(+), 11 deletions(-) create mode 100644 tests/code_coverage_tests/check_py310_typing_imports.py create mode 100644 tests/test_litellm/test_check_py310_typing_imports.py diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index c112bf2bb22..d02f5878396 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -128,6 +128,9 @@ jobs: - name: check_fastuuid_usage run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py + - name: check_py310_typing_imports + run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py + - name: check_e2e_no_raw_requests run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -145,3 +148,33 @@ jobs: - name: documentation_test_api_docs run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py + + python-310-import-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.10" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: uv sync --frozen --extra proxy --python 3.10 + + - run: uv run --no-sync python --version + + - name: Import litellm + run: uv run --no-sync python -c "import litellm" + + - name: Check litellm CLI + run: uv run --no-sync litellm --version diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2d737bc34e7..587da997f94 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,9 +10,9 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast -from typing_extensions import ReadOnly +from typing_extensions import Never, ReadOnly import litellm from litellm._logging import verbose_logger diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 4f6305d88cf..af02c11ad86 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -5,10 +5,10 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never +from typing import TYPE_CHECKING, Any, Final, TypedDict from pydantic import ValidationError -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 3e4dc07a521..3b8151d1064 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -13,10 +13,10 @@ import os import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Annotated, Final, TypedDict, assert_never +from typing import Annotated, Final, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1682cf12f4e..47f69732e95 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,7 +7,9 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar, assert_never +from typing import Final, Literal, Protocol, TypeVar + +from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 0aee5e8cc54..a41bd36d510 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -2,7 +2,9 @@ import json from collections.abc import Iterator from dataclasses import dataclass from itertools import chain -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 9d450cb5b8d..8c59a520272 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose. from dataclasses import dataclass from pathlib import Path -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.path_utils import safe_filename diff --git a/litellm/types/llms/gemini_audio_transcription.py b/litellm/types/llms/gemini_audio_transcription.py index cb12e0f45b8..f7e74ba4bf8 100644 --- a/litellm/types/llms/gemini_audio_transcription.py +++ b/litellm/types/llms/gemini_audio_transcription.py @@ -1,7 +1,7 @@ -from typing import Literal, Required +from typing import Literal from pydantic import BaseModel, ConfigDict -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict class GeminiTranscriptionAudioInput(TypedDict): diff --git a/tests/code_coverage_tests/check_py310_typing_imports.py b/tests/code_coverage_tests/check_py310_typing_imports.py new file mode 100644 index 00000000000..0cd4d089890 --- /dev/null +++ b/tests/code_coverage_tests/check_py310_typing_imports.py @@ -0,0 +1,150 @@ +import ast +import os +import sys +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset( + { + "NotRequired", + "Required", + "Self", + "LiteralString", + "Never", + "assert_never", + "assert_type", + "reveal_type", + "TypeVarTuple", + "Unpack", + "dataclass_transform", + "override", + "TypeAliasType", + "get_original_bases", + "ReadOnly", + "TypeIs", + "NoDefault", + "get_protocol_members", + "is_protocol", + "evaluate_forward_ref", + "TypeForm", + } +) + + +@dataclass(frozen=True, slots=True) +class TypingImportViolation: + file: str + line: int + name: str + + +def _walk_with_ancestors( + node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = () +) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]: + yield node, ancestors + for field_name, field_value in ast.iter_fields(node): + if isinstance(field_value, ast.AST): + yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name))) + elif isinstance(field_value, list): + for child in field_value: + if isinstance(child, ast.AST): + yield from _walk_with_ancestors(child, (*ancestors, (node, field_name))) + + +def _is_sys_version_info(node: ast.AST) -> bool: + return ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "sys" + and node.attr == "version_info" + ) + + +def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool: + nearest_if: Final[tuple[ast.If, str] | None] = next( + ( + (ancestor, field_name) + for ancestor, field_name in reversed(ancestors) + if isinstance(ancestor, ast.If) + ), + None, + ) + if nearest_if is None: + return False + enclosing_if, branch = nearest_if + test: Final[ast.expr] = enclosing_if.test + if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left): + return False + operator: Final[ast.cmpop] = test.ops[0] + return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or ( + isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse" + ) + + +def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]: + path: Final[Path] = Path(file_path) + tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + return tuple( + violation + for node, ancestors in _walk_with_ancestors(tree) + if not _is_version_guarded(ancestors) + for violation in _violations_for_node(node, path) + ) + + +def _violations_for_node( + node: ast.AST, path: Path +) -> tuple[TypingImportViolation, ...]: + if isinstance(node, ast.ImportFrom) and node.module == "typing": + return tuple( + TypingImportViolation(file=str(path), line=node.lineno, name=alias.name) + for alias in node.names + if alias.name in PY311_PLUS_TYPING_NAMES + ) + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "typing" + and node.attr in PY311_PLUS_TYPING_NAMES + ): + return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),) + return () + + +def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]: + base_path: Final[Path] = Path(base_dir) + return tuple( + violation + for directory in ( + base_path / "litellm", + base_path / "enterprise", + base_path / "litellm-proxy-extras" / "litellm_proxy_extras", + ) + if directory.exists() + for path in directory.rglob("*.py") + for violation in scan_file(path) + ) + + +def main() -> None: + violations: Final[tuple[TypingImportViolation, ...]] = scan_directory() + if violations: + message: Final[str] = "\n".join( + ( + "Python 3.10-incompatible typing imports found:", + *( + f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; " + "import it from typing_extensions instead because litellm supports Python 3.10" + for violation in violations + ), + ) + ) + sys.stdout.write(f"{message}\n") + raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead") + sys.stdout.write("No Python 3.10-incompatible typing imports found.\n") + + +if __name__ == "__main__": + main() diff --git a/tests/test_litellm/test_check_py310_typing_imports.py b/tests/test_litellm/test_check_py310_typing_imports.py new file mode 100644 index 00000000000..de370326091 --- /dev/null +++ b/tests/test_litellm/test_check_py310_typing_imports.py @@ -0,0 +1,86 @@ +import sys +from pathlib import Path +from typing import Final + +_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests" +sys.path.insert(0, str(_CODE_COVERAGE_DIR)) # test-quality-ok: required to import checker from its source directory +import check_py310_typing_imports as checker # noqa: E402 # load checker from its source directory + + +def _scan(tmp_path: Path, source: str) -> tuple[object, ...]: + file_path = tmp_path / "fixture.py" + file_path.write_text(source, encoding="utf-8") + return checker.scan_file(file_path) + + +def test_typing_import_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n") + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_typing_extensions_import_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == () + + +def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "import typing\nx: typing.Self\n") + assert tuple(violation.name for violation in violations) == ("Self",) + + +def test_version_guarded_typing_import_passes(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing import NotRequired\n" + "else:\n" + " from typing_extensions import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + violations = _scan(tmp_path, source) + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info < (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_nearest_if_controls_version_guard(tmp_path: Path) -> None: + source = ( + "if sys.version_info >= (3, 11):\n" + " from typing import Self\n" + " x = 1\n" + "if True:\n" + " from typing import Self\n" + ) + violations = _scan(tmp_path, source) + assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),) + + +def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None: + file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py" + file_path.parent.mkdir(parents=True) + file_path.write_text("from typing import NotRequired\n", encoding="utf-8") + + violations = checker.scan_directory(tmp_path) + assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),) + + +def test_python_310_typing_name_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing import Optional\n") == () From 8441dd6e8ca16a480e0766651f7c4ab99d22d5d5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:28:06 -0700 Subject: [PATCH 6/6] fix(proxy): keep SpendLogs and callback session ids in sync when the request has none (#39450) * fix(proxy): keep SpendLogs and callback session ids in sync when the request has none Add general_settings.missing_session_id (generate | reject). In generate mode one id is stamped into litellm_session_id, litellm_trace_id and metadata.session_id before callbacks run, so LiteLLM_SpendLogs.session_id and the Langfuse session id match. In reject mode such requests get a 400. Unset keeps the legacy behavior. MCP routes are not affected Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(proxy): regenerate schema.d.ts and shorten mutable-ok comment for ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark generated session ids so affinity consumers do not pin on them Fireworks x-session-affinity, the router session_affinity pre-call check and the complexity router session pin all read metadata.session_id as a caller-chosen stable key. A missing_session_id: generate id is fresh per request, so it now carries metadata.litellm_session_id_generated and those consumers skip it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/llms/fireworks_ai/common_utils.py | 11 +- litellm/proxy/_types.py | 4 + litellm/proxy/litellm_pre_call_utils.py | 55 ++++++ .../complexity_router/complexity_router.py | 8 +- .../deployment_affinity_check.py | 4 +- .../test_fireworks_ai_chat_transformation.py | 16 ++ .../proxy/test_litellm_pre_call_utils.py | 176 ++++++++++++++++++ .../router_strategy/test_complexity_router.py | 22 ++- .../test_session_id_affinity.py | 43 ++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 11 files changed, 335 insertions(+), 10 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c7b74e176db..ef9329b9dfc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 8c306faa036..ac934ad0cb5 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -2,6 +2,7 @@ from typing import Final from httpx import Headers +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. - Deliberately does not fall back to `litellm_trace_id`: that is generated per - request (`str(uuid.uuid4())` when absent), so using it pins every request to a - different Fireworks node and prompt caching never hits. + Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the + proxy generated for a request that had none: both are per request, so using them + pins every request to a different Fireworks node and prompt caching never hits. """ params: Final = litellm_params + metadata: Final = params.get("metadata") + if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - metadata: Final = params.get("metadata") if isinstance(metadata, dict): value = metadata.get("session_id") if value: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2da7ceb2d50..849e54c65aa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", ) + missing_session_id: Literal["generate", "reject"] | None = Field( + None, + description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", + ) enable_public_model_hub: bool = Field( default=False, description="Public model hub for users to see what models they have access to, supported openai params, etc.", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 20f83085286..1d440448c2f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -16,6 +16,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm._uuid import uuid from litellm.constants import ( CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -23,6 +24,7 @@ from litellm.constants import ( OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -40,6 +42,7 @@ from litellm.proxy._types import ( AddTeamCallback, CommonProxyErrors, LitellmDataForBackendLLMCall, + LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -47,6 +50,8 @@ from litellm.proxy._types import ( TeamCallbackMetadata, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import get_request_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, @@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: return session_id +def _is_llm_inference_route(request: Request) -> bool: + route: Final = get_request_route(request) + return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) + + +def apply_missing_session_id_policy( + data: dict[str, object], # mutable-ok: stamps session ids in place on the request body the pipeline threads through + _metadata_variable_name: str, + general_settings: Mapping[str, object] | None, + request: Request, +) -> None: + policy: Final = general_settings.get("missing_session_id") if general_settings else None + if policy is None or not _is_llm_inference_route(request): + return + metadata: Final = data.get(_metadata_variable_name) + if not isinstance(metadata, dict): + return + if data.get("litellm_session_id") or metadata.get("session_id"): + return + match policy: + case "generate": + session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4()) + data["litellm_session_id"] = session_id # rebind-ok: data is an out-param + data.setdefault("litellm_trace_id", session_id) + metadata["session_id"] = session_id + metadata[SESSION_ID_GENERATED_METADATA_KEY] = True + case "reject": + raise ProxyException( + message=( + "Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. " + "Required by `general_settings.missing_session_id: reject`." + ), + type=ProxyErrorTypes.bad_request_error, + param="session_id", + code=400, + ) + case _: + verbose_proxy_logger.warning( + "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy + ) + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code identifies itself as ``claude-cli/ ...``; the IDE extensions and the Agent SDK run through the same CLI and share that prefix.""" @@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request( data=data, _metadata_variable_name=_metadata_variable_name, ) + apply_missing_session_id_policy( + data=data, + _metadata_variable_name=_metadata_variable_name, + general_settings=general_settings, + request=request, + ) # Expose request headers under the metadata field for guardrails (fixes #17477) if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict): diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index d205db90607..430efe339a2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, create_model from litellm._logging import verbose_router_logger -from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import ( + EMPTY_MAPPING, + RETURN_RAW_MODEL_NAME_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata @@ -2712,7 +2716,7 @@ class ComplexityRouter(CustomLogger): """Resolve a client-supplied session_id.""" for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): session_id = metadata.get("session_id") - if session_id is not None: + if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return str(session_id) return None diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b1e9dbdefa8..39d3e25aacb 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -21,7 +21,7 @@ from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") - if session_id is None: + if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None return str(session_id) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d7cc89868af..ec8725db5f7 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -8,6 +8,7 @@ import litellm from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): ) +def test_get_fireworks_session_id_ignores_proxy_generated_session_id(): + """general_settings.missing_session_id: generate stamps a fresh id per request; sending it + as x-session-affinity would pin every request to a different node.""" + assert ( + get_fireworks_session_id( + { + "litellm_session_id": "generated-1", + "litellm_trace_id": "generated-1", + "metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + } + ) + is None + ) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( 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 ee0e2014951..8366e5546a9 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem @@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): } assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + + +def _request_for(path: str) -> MagicMock: + request = MagicMock(spec=Request) + request.scope = {"path": path} + request.url = MagicMock() + request.url.path = path + request.url.__str__.return_value = f"http://localhost{path}" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + return request + + +def _spend_log_session_id(data: dict[str, object]) -> str: + """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log + + metadata = data["metadata"] + assert isinstance(metadata, dict) + litellm_params = get_litellm_params( + litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None, + litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None, + metadata=metadata, + ) + trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"), + litellm_params=litellm_params, + ) + return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_correlation_in_logs", [False, True]) +async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree( + monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool +): + """Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs + must be the same generated id, so cross-referencing the two by session_id works. The id is marked + as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs) + data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} + + updated = await add_litellm_data_to_request( + data=data, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + callback_session_id = updated["metadata"]["session_id"] + assert isinstance(callback_session_id, str) and len(callback_session_id) == 36 + assert _spend_log_session_id(updated) == callback_session_id + assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True + assert get_fireworks_session_id( + {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]} + ) is None + + +@pytest.mark.asyncio +async def test_missing_session_id_unset_keeps_legacy_divergence(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + + assert "session_id" not in updated["metadata"] + assert "litellm_session_id" not in updated + assert _spend_log_session_id(updated) == "per-call-random-trace-id" + + +@pytest.mark.asyncio +async def test_missing_session_id_generate_reuses_traceparent_trace_id(): + """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" + request = _request_for("/v1/chat/completions") + request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["generate", "reject"]) +async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str): + request = _request_for("/v1/chat/completions") + request.headers = {"x-litellm-session-id": "client-session-1"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": policy}, + ) + + assert updated["litellm_session_id"] == "client-session-1" + assert updated["metadata"]["session_id"] == "client-session-1" + assert _spend_log_session_id(updated) == "client-session-1" + assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"] + assert ( + get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]}) + == "client-session-1" + ) + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_accepts_body_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert updated["metadata"]["session_id"] == "body-session-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_returns_400_without_session_id(): + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "session_id" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"]) +async def test_missing_session_id_policy_skips_non_inference_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert "session_id" not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_missing_session_id_unknown_value_is_ignored(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "typo"}, + ) + + assert "session_id" not in updated["metadata"] diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 941d78085e4..d7d02544efb 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -16,7 +16,7 @@ import litellm from litellm import Router from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -4274,6 +4274,26 @@ class TestSessionAffinity: assert first.model == "o1-preview" assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio + async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config): + """A session id the proxy generated for a request that had none is per request, so + it must not create a pin even with session_affinity enabled.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}} + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index a3772a276fa..cf48888600e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -7,7 +7,7 @@ import json import litellm from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert filtered[0]["model_info"]["id"] == "deployment-2" +@pytest.mark.asyncio +async def test_proxy_generated_session_id_does_not_pin_a_deployment(): + """A session id the proxy generated for a request that had none is per request, so a + pin stored under it must be ignored and none must be written.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + healthy_deployments = [ + {"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}}, + {"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}}, + ] + await cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"), + {"model_id": "deployment-2"}, + ) + request_kwargs = { + "metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True} + } + + filtered = await callback.async_filter_deployments( + model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs + ) + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"}, + "model_info": {"id": "deployment-1"}, + }, + call_type=None, + ) + + assert len(filtered) == 2 + assert await cache.async_get_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1") + ) == {"model_id": "deployment-2"} + + MOCK_RESPONSES_API_RESPONSE = { "id": "resp_mock-resp-456", "object": "response", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 90c4f03bf08..491c4fb6a44 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25772,6 +25772,11 @@ export interface components { * @description Number of trusted reverse proxies/load balancers in front of the gateway that append to X-Forwarded-For. When set (and mcp_trusted_proxy_ranges validates the direct peer), the client IP for MCP access control is read this many entries from the right of the chain instead of the spoofable leftmost value, defeating append-style X-Forwarded-For forgery. */ mcp_xff_num_trusted_hops?: number | null; + /** + * Missing Session Id + * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. + */ + missing_session_id?: ("generate" | "reject") | null; /** * Model List Healthy Only * @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called.