From 698f608ad6ee78ee7d84d8947c386244ea7b1e4e Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 18:26:15 -0700 Subject: [PATCH 001/297] fix(guardrails): transfer guardrail evaluation metadata to spend logs on success path On successful requests, guardrail evaluations run and store results in request_data["litellm_metadata"]["standard_logging_guardrail_information"], but the spend-log serialization reads from request_data["metadata"], so guardrail evaluations were never reported in Request Logs. The failure path already had this transfer; this adds the same logic to the success path so guardrail evaluation info appears in both success and failure spend logs. Fixes LIT-6314. --- .../proxy/hooks/proxy_track_cost_callback.py | 11 ++ .../hooks/test_proxy_track_cost_callback.py | 132 ++++++++++-------- 2 files changed, 85 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 47aafda2337..ca8341822e9 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -269,6 +269,17 @@ class _ProxyDBLogger(CustomLogger): served_model_id=sl_object.get("model_id") if sl_object is not None else None, router=get_llm_router(), ) + # LIT-6314: post-call guardrail evaluations run after SLP is built, so + # guardrail_information is missing from SLP. Populate it from post-call evals + # before spend-log write so reports show accurate guardrail results. + if sl_object is not None: + guardrail_info_from_hooks: Final = ( + kwargs.get("litellm_metadata", {}).get("standard_logging_guardrail_information") + if isinstance(kwargs.get("litellm_metadata"), dict) + else None + ) + if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"): + sl_object["guardrail_information"] = guardrail_info_from_hooks if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 8043a1aca3f..5eb59acf4f5 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,4 +1,3 @@ - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -70,9 +69,7 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { - "request_id": "test_request_id" - } + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -336,9 +333,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs[ - "budget_reservation" - ] + mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -433,36 +428,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } + assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} - ) - is None - ) - assert ( - _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": UserAPIKeyAuth( - budget_reservation=budget_reservation - ) - } + metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": dict( - UserAPIKeyAuth(budget_reservation=budget_reservation) - ) - } + metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata( - metadata={"user_api_key_budget_reservation": budget_reservation} - ) + _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) is budget_reservation ) @@ -470,9 +450,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=Exception("db unavailable") - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -508,9 +486,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=db_exception - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -554,12 +530,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call( - "Failed to release budget reservation after database update failed" - ) - mock_log_exception.assert_any_call( - "Failed to invalidate budget reservation counters after release failed" - ) + mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") + mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") increment_spend_counters.assert_not_awaited() @@ -1101,10 +1073,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert ( - call_kwargs["standard_logging_object"]["trace_id"] - == "trace-id-from-logging-obj" - ) + assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1691,9 +1660,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } with patch( @@ -1772,15 +1739,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs - ) + update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] - == "mcp-user@example.com" - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" @pytest.mark.parametrize( @@ -1828,9 +1790,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request( - call_type, expect_spend_log -): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1876,9 +1836,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( - 1 if expect_spend_log else 0 - ) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) class _FakeDeploymentLookup: @@ -1989,3 +1947,61 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un ) assert charged == ("premium", "tier0") + + +@pytest.mark.asyncio +async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object(): + """ + LIT-6314 regression: post-call guardrails append evaluations to + request_data["litellm_metadata"]["standard_logging_guardrail_information"], + but the standard_logging_object (built pre-post-call) lacks this field. + The callback must populate standard_logging_object["guardrail_information"] + from post-call evals before spend-log write so reports show results. + """ + logger = _ProxyDBLogger() + guardrail_info = [ + { + "guardrail_name": "bedrock-guard", + "guardrail_status": "success", + "guardrail_cost": 0.0001, + } + ] + sl_object = {"response_cost": 0.01} + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, + "litellm_params": { + "metadata": {"user_api_key": "test_key"}, + }, + "call_type": CallTypes.completion.value, + "response_cost": 0.01, + "standard_logging_object": sl_object, + } + + with ( + patch( + "litellm.proxy.proxy_server.increment_spend_counters", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.update_cache", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.failed_tracking_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == 1 + assert sl_object.get("guardrail_information") == guardrail_info From ca9e121eb67185debdf1af3dea64ab2e080f3df1 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 18:57:20 -0700 Subject: [PATCH 002/297] fix(guardrails): avoid mutable dict literal in guardrail metadata population --- litellm/proxy/hooks/proxy_track_cost_callback.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ca8341822e9..b78884f7d95 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -273,13 +273,16 @@ class _ProxyDBLogger(CustomLogger): # guardrail_information is missing from SLP. Populate it from post-call evals # before spend-log write so reports show accurate guardrail results. if sl_object is not None: + litellm_metadata: Final = kwargs.get("litellm_metadata") guardrail_info_from_hooks: Final = ( - kwargs.get("litellm_metadata", {}).get("standard_logging_guardrail_information") - if isinstance(kwargs.get("litellm_metadata"), dict) + litellm_metadata.get("standard_logging_guardrail_information") + if isinstance(litellm_metadata, dict) else None ) if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"): - sl_object["guardrail_information"] = guardrail_info_from_hooks + sl_object["guardrail_information"] = ( + guardrail_info_from_hooks # mutable-ok: populate SLP before spend-log write + ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) From 686e8ddae7aaac16089de238d09c676b6b62824a Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 19:04:14 -0700 Subject: [PATCH 003/297] test: suppress TQ008 on proxy_server global patches, matching file idiom --- .../proxy/hooks/test_proxy_track_cost_callback.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 5eb59acf4f5..2d5857f95e5 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1980,15 +1980,15 @@ async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object(): } with ( - patch( + patch( # test-quality-ok: spend counters are a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock, ), - patch( + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock, ), - patch( + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam "litellm.proxy.proxy_server.proxy_logging_obj", ) as mock_proxy_logging, ): From 0e6c4279b6387a26ee43e7ef8b8454bdee65295e Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 21:35:32 -0700 Subject: [PATCH 004/297] fix(guardrails): record not_run evaluation when scoping leaves nothing to scan Replaces the metadata transfer approach: that block read top-level litellm_metadata which guardrail info never populates, and the SLP builder already reads the nested bucket it lands in, so it was dead code and is reverted. Real cause of missing evaluations: process_input_messages skips apply_guardrail entirely when message scoping (skip_system_message, skip_tool, scan_only_tool_results) leaves no scannable content, so the guardrail shows up in applied_guardrails with no guardrail_information entry. Now records a not_run entry unless the guardrail records its own information. --- .../chat/guardrail_translation/handler.py | 9 ++ .../proxy/hooks/proxy_track_cost_callback.py | 14 -- .../test_openai_guardrail_handler.py | 50 +++++++ .../hooks/test_proxy_track_cost_callback.py | 132 ++++++++---------- 4 files changed, 117 insertions(+), 88 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index de15fefe943..76a7f33454a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -191,6 +191,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) + elif not guardrail_to_apply.records_own_guardrail_information: + # Every guardrail in applied_guardrails needs a persisted evaluation record, + # or request logs report it as silently missing (LIT-6314). + guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="no scannable content after message scoping", + request_data=data, + guardrail_status="not_run", + ) + verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed input messages: %s", data.get("messages"), diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b78884f7d95..47aafda2337 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -269,20 +269,6 @@ class _ProxyDBLogger(CustomLogger): served_model_id=sl_object.get("model_id") if sl_object is not None else None, router=get_llm_router(), ) - # LIT-6314: post-call guardrail evaluations run after SLP is built, so - # guardrail_information is missing from SLP. Populate it from post-call evals - # before spend-log write so reports show accurate guardrail results. - if sl_object is not None: - litellm_metadata: Final = kwargs.get("litellm_metadata") - guardrail_info_from_hooks: Final = ( - litellm_metadata.get("standard_logging_guardrail_information") - if isinstance(litellm_metadata, dict) - else None - ) - if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"): - sl_object["guardrail_information"] = ( - guardrail_info_from_hooks # mutable-ok: populate SLP before spend-log write - ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a29e0be4655..24060e94e54 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 @@ -1559,3 +1559,53 @@ class TestScanOnlyToolResults: assert data["messages"][3]["content"] == "page says [BLOCKED] here" assert data["messages"][3]["tool_call_id"] == "call_1" assert data["messages"][4]["content"] == "and then?" + + +class TestNoScannableContentRecordsNotRun: + """LIT-6314: a guardrail whose scoping leaves nothing to scan must still persist an evaluation record""" + + def _system_only_data(self) -> dict: + return {"messages": [{"role": "system", "content": "SYSTEM-PROMPT"}]} + + def _recorded_entries(self, data: dict) -> list: + metadata = data.get("metadata") or data.get("litellm_metadata") or {} + return metadata.get("standard_logging_guardrail_information") or [] + + @pytest.mark.asyncio + async def test_skipped_scan_records_not_run_entry(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="skip-system-guardrail") + guardrail.skip_system_message_in_guardrail = True + data = self._system_only_data() + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None, "nothing survived scoping, apply_guardrail must not run" + entries = self._recorded_entries(data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "skip-system-guardrail" + assert entries[0]["guardrail_status"] == "not_run" + + @pytest.mark.asyncio + async def test_self_recording_guardrail_is_left_alone(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="self-recording-guardrail") + guardrail.skip_system_message_in_guardrail = True + guardrail.records_own_guardrail_information = True + data = self._system_only_data() + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None + assert self._recorded_entries(data) == [] + + @pytest.mark.asyncio + async def test_scannable_content_records_no_extra_entry(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="normal-guardrail") + data = {"messages": [{"role": "user", "content": "hello"}]} + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is not None + assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data)) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 2d5857f95e5..8043a1aca3f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,3 +1,4 @@ + from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -69,7 +70,9 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { + "request_id": "test_request_id" + } metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -333,7 +336,9 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] + mock_invalidate_budget_reservation_counters.await_args.kwargs[ + "budget_reservation" + ] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -428,21 +433,36 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } - assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} + metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} + ) + is None + ) + assert ( + _get_budget_reservation_from_metadata( + metadata={ + "user_api_key_auth": UserAPIKeyAuth( + budget_reservation=budget_reservation + ) + } ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} + metadata={ + "user_api_key_auth": dict( + UserAPIKeyAuth(budget_reservation=budget_reservation) + ) + } ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) + _get_budget_reservation_from_metadata( + metadata={"user_api_key_budget_reservation": budget_reservation} + ) is budget_reservation ) @@ -450,7 +470,9 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + side_effect=Exception("db unavailable") + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -486,7 +508,9 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + side_effect=db_exception + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -530,8 +554,12 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") - mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") + mock_log_exception.assert_any_call( + "Failed to release budget reservation after database update failed" + ) + mock_log_exception.assert_any_call( + "Failed to invalidate budget reservation counters after release failed" + ) increment_spend_counters.assert_not_awaited() @@ -1073,7 +1101,10 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" + assert ( + call_kwargs["standard_logging_object"]["trace_id"] + == "trace-id-from-logging-obj" + ) # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1660,7 +1691,9 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + "combined_usage_object": Usage( + prompt_tokens=30, completion_tokens=1, total_tokens=31 + ), } with patch( @@ -1739,10 +1772,15 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs + update_kwargs = ( + mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs + ) assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" + assert ( + kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] + == "mcp-user@example.com" + ) @pytest.mark.parametrize( @@ -1790,7 +1828,9 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request( + call_type, expect_spend_log +): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1836,7 +1876,9 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(cal end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( + 1 if expect_spend_log else 0 + ) class _FakeDeploymentLookup: @@ -1947,61 +1989,3 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un ) assert charged == ("premium", "tier0") - - -@pytest.mark.asyncio -async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object(): - """ - LIT-6314 regression: post-call guardrails append evaluations to - request_data["litellm_metadata"]["standard_logging_guardrail_information"], - but the standard_logging_object (built pre-post-call) lacks this field. - The callback must populate standard_logging_object["guardrail_information"] - from post-call evals before spend-log write so reports show results. - """ - logger = _ProxyDBLogger() - guardrail_info = [ - { - "guardrail_name": "bedrock-guard", - "guardrail_status": "success", - "guardrail_cost": 0.0001, - } - ] - sl_object = {"response_cost": 0.01} - kwargs = { - "model": "gpt-4", - "messages": [{"role": "user", "content": "Hello"}], - "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, - "litellm_params": { - "metadata": {"user_api_key": "test_key"}, - }, - "call_type": CallTypes.completion.value, - "response_cost": 0.01, - "standard_logging_object": sl_object, - } - - with ( - patch( # test-quality-ok: spend counters are a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.increment_spend_counters", - new_callable=AsyncMock, - ), - patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.update_cache", - new_callable=AsyncMock, - ), - patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam - "litellm.proxy.proxy_server.proxy_logging_obj", - ) as mock_proxy_logging, - ): - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() - mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() - mock_proxy_logging.failed_tracking_alert = AsyncMock() - - await logger._PROXY_track_cost_callback( - kwargs=kwargs, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == 1 - assert sl_object.get("guardrail_information") == guardrail_info From 527c36343fd93ebaeffec0d4e07f6f9959f85da9 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 00:20:14 -0700 Subject: [PATCH 005/297] fix(guardrails): keep not_run entries out of daily evaluation counts --- litellm/proxy/guardrails/usage_tracking.py | 21 +++++++------ .../proxy/guardrails/test_usage_tracking.py | 31 +++++++++++++++++++ 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index b8ae09afc00..ec059acb146 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -316,15 +316,18 @@ async def process_spend_logs_guardrail_usage( guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" if not guardrail_id: continue - key = _MetricsKey(guardrail_id, date_key) - daily_guardrail[key]["requests_evaluated"] += 1 - action = _guardrail_status_to_action(entry.get("guardrail_status")) - if action == "passed": - daily_guardrail[key]["passed_count"] += 1 - elif action == "blocked": - daily_guardrail[key]["blocked_count"] += 1 - else: - daily_guardrail[key]["flagged_count"] += 1 + status = entry.get("guardrail_status") + # not_run means the guardrail never evaluated the request: index it for drill-down, keep it out of counts + if status != "not_run": + key = _MetricsKey(guardrail_id, date_key) + daily_guardrail[key]["requests_evaluated"] += 1 + action = _guardrail_status_to_action(status) + if action == "passed": + daily_guardrail[key]["passed_count"] += 1 + elif action == "blocked": + daily_guardrail[key]["blocked_count"] += 1 + else: + daily_guardrail[key]["flagged_count"] += 1 policy_id = entry.get("policy_id") index_rows.append( { diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 6da121703d7..b98b037e7b8 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -306,6 +306,37 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): } +@pytest.mark.asyncio +async def test_not_run_entries_are_indexed_but_not_counted_as_evaluations(): + """ + LIT-6314 records a not_run entry when message scoping leaves a guardrail + nothing to scan. The guardrail never evaluated the request, so counting it + as a passed evaluation would inflate daily pass rates; it still gets an + index row so per-request drill-down finds the spend log. + """ + prisma = _prisma() + logs = [_payload("r1", guardrail_status="not_run"), _payload("r2")] + + await process_spend_logs_guardrail_usage(prisma, logs) + + metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] + assert metrics_create["requests_evaluated"] == 1 + assert metrics_create["passed_count"] == 1 + index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] + assert sorted(row["request_id"] for row in index_rows) == ["r1", "r2"] + + +@pytest.mark.asyncio +async def test_batch_of_only_not_run_entries_writes_no_metrics_row(): + prisma = _prisma() + + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="not_run")]) + + assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 0 + index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] + assert [row["request_id"] for row in index_rows] == ["r1"] + + @pytest.mark.asyncio async def test_payload_without_request_id_is_skipped_like_the_metrics_path(): prisma = _prisma() From 45ff658c6f020665bbe997ba6fc93ebf7c4eb2e3 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 01:09:01 -0700 Subject: [PATCH 006/297] fix(ui): render not_run guardrail evaluations as not run instead of failed --- ui/litellm-dashboard/eslint-suppressions.json | 2 +- .../GuardrailViewer/GuardrailViewer.test.tsx | 10 +++ .../GuardrailViewer/GuardrailViewer.tsx | 71 +++++++++++++------ 3 files changed, 60 insertions(+), 23 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7de7373b20b..93966c8d0ee 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2352,7 +2352,7 @@ }, "src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx": { "no-nested-ternary": { - "count": 4 + "count": 3 } }, "src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx": { diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index b5e04c72440..0fb5504d231 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -33,6 +33,16 @@ describe("GuardrailViewer", () => { expect(screen.getByText("1235ms")).toBeInTheDocument(); }); + it("renders not_run entries as not run instead of failed", () => { + const data = makeGuardrailInformation({ guardrail_status: "not_run", guardrail_mode: "pre_call" }); + renderWithProviders(); + + expect(screen.getByText(/1 Not run/)).toBeInTheDocument(); + // one NOT RUN badge in the timeline, one in the evaluation card + expect(screen.getAllByText("NOT RUN")).toHaveLength(2); + expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); + }); + it("calculates and displays masked entity totals", async () => { const user = userEvent.setup(); const data = makeGuardrailInformation({ diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index 863f4117510..08f9002b050 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -137,6 +137,36 @@ const isEntrySuccess = (entry: GuardrailInformation): boolean => { return (entry.guardrail_status ?? "").toLowerCase() === "success"; }; +const isEntryNotRun = (entry: GuardrailInformation): boolean => { + return (entry.guardrail_status ?? "").toLowerCase() === "not_run"; +}; + +type EntryStatusLabel = "PASSED" | "NOT RUN" | "FAILED"; + +const entryStatusLabel = (entry: GuardrailInformation): EntryStatusLabel => { + if (isEntrySuccess(entry)) return "PASSED"; + if (isEntryNotRun(entry)) return "NOT RUN"; + return "FAILED"; +}; + +const StatusIcon = ({ status }: { status: EntryStatusLabel }) => { + if (status === "PASSED") return ; + if (status === "NOT RUN") return ; + return ; +}; + +const timelineStatusClass = (status: EntryStatusLabel): string => { + if (status === "PASSED") return "bg-success/15 text-success"; + if (status === "NOT RUN") return "bg-muted text-muted-foreground"; + return "bg-destructive/15 text-destructive"; +}; + +const cardStatusClass = (status: EntryStatusLabel): string => { + if (status === "PASSED") return "bg-success/15 text-success border border-success/20"; + if (status === "NOT RUN") return "bg-muted text-muted-foreground border border-border"; + return "bg-destructive/15 text-destructive border border-destructive/20"; +}; + const getRiskColor = (score: number): string => { if (score <= 3) return "text-success bg-success/10 border-success/20"; if (score <= 6) return "text-warning bg-warning/10 border-warning/20"; @@ -318,8 +348,7 @@ interface TimelineEntry { type: "request" | "guardrail" | "llm" | "response"; label: string; offsetMs: number; - status?: string; - isSuccess?: boolean; + status?: EntryStatusLabel; } const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { @@ -348,8 +377,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `Pre-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + status: entryStatusLabel(e), }); } @@ -372,8 +400,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `During-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + status: entryStatusLabel(e), }); } @@ -384,8 +411,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `Post-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + status: entryStatusLabel(e), }); } @@ -410,10 +436,8 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { ) : item.type === "llm" ? ( - ) : item.isSuccess ? ( - ) : ( - + )} {idx < timeline.length - 1 &&
} @@ -427,9 +451,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { {item.status && ( {item.status} @@ -456,6 +478,7 @@ const formatGuardrailCost = (cost: number): string => { const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { const [expanded, setExpanded] = useState(false); const success = isEntrySuccess(entry); + const statusLabel = entryStatusLabel(entry); const totalMasked = getTotalMasked(entry); const displayName = getDisplayName(entry); const durationStr = formatDurationMs(entry.duration); @@ -490,7 +513,9 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { onClick={() => setExpanded(!expanded)} > {/* Status icon */} -
{success ? : }
+
+ +
{/* Name + badges */}
@@ -501,13 +526,9 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { - {success ? "PASSED" : "FAILED"} + {statusLabel} {matchCountStr && ( @@ -673,7 +694,8 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) }, [data]); const passedCount = guardrailEntries.filter(isEntrySuccess).length; - const allPassed = passedCount === guardrailEntries.length; + const notRunCount = guardrailEntries.filter(isEntryNotRun).length; + const allPassed = passedCount === guardrailEntries.length - notRunCount; const totalOverheadMs = useMemo(() => { return Math.round(guardrailEntries.reduce((sum, e) => sum + (e.duration ?? 0), 0) * 1000); @@ -728,6 +750,11 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) ) : null} {passedCount} Passed + {notRunCount > 0 && ( + + {notRunCount} Not run + + )}
From 26fc1ff221ba445aadf39fa10ead22d759761afe Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 01:09:01 -0700 Subject: [PATCH 007/297] fix(guardrails): stop readers from scoring not_run evaluations as passed --- litellm/proxy/compliance_checks.py | 3 +- litellm/proxy/guardrails/usage_endpoints.py | 6 +- .../proxy/guardrails/test_usage_endpoints.py | 70 +++++++++++++++++-- .../test_compliance_endpoints.py | 36 +++++++++- .../GuardrailsMonitor/LogViewer.tsx | 11 ++- .../components/GuardrailsMonitor/mockData.ts | 2 +- .../LogDetailContent.test.tsx | 29 +++++++- .../LogDetailsDrawer/LogDetailContent.tsx | 4 +- 8 files changed, 147 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index ff311911742..053c88d10ed 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -26,7 +26,8 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = data.guardrail_information or [] + # a not_run entry records a guardrail that never evaluated the request, so it cannot evidence compliance + self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "not_run"] def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 7a0edbddca8..75ffc1545dc 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -256,7 +256,7 @@ class UsageDetailResponse(BaseModel): class UsageLogEntry(BaseModel): id: str timestamp: str - action: str # blocked | passed | flagged + action: str # blocked | passed | flagged | not_run score: float | None latency_ms: float | None model: str | None @@ -691,7 +691,9 @@ def _usage_log_entry_from_row( reason_val = None if entry_for_guardrail: st: Final = (entry_for_guardrail.get("guardrail_status") or "").lower() - if "intervened" in st or "block" in st: + if st == "not_run": + action_val = "not_run" + elif "intervened" in st or "block" in st: action_val = "blocked" elif "fail" in st or "error" in st: action_val = "flagged" diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index 1665fa03639..e469aef9a61 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -9,12 +9,10 @@ orphans), and logs missed their logical-name alias. """ from datetime import datetime -from typing import Any, Optional +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest - - from fastapi import HTTPException from prisma.errors import TableNotFoundError @@ -47,7 +45,7 @@ def _yaml_guardrail( guardrail_id: str = "yaml-1", name: str = "yaml-pii", provider: str = "presidio", - info: Optional[dict] = None, + info: dict | None = None, ) -> Guardrail: return Guardrail( guardrail_id=guardrail_id, @@ -433,3 +431,67 @@ async def test_detail_prev_trend_query_is_bounded(): prev_wheres = [w for w in wheres if "lt" in w.get("date", {})] assert prev_wheres assert all("gte" in w["date"] for w in prev_wheres) + + +@pytest.mark.asyncio +async def test_logs_report_not_run_entries_as_not_run_not_passed(): + """LIT-6314: a guardrail that never scanned must not be reported as a pass in the drill-down.""" + index_row = MagicMock() + index_row.request_id = "req-nr" + index_row.guardrail_id = "db-1" + index_row.start_time = datetime(2026, 4, 22) + spend_log = MagicMock() + spend_log.request_id = "req-nr" + spend_log.model = "gpt-4o-mini" + spend_log.startTime = datetime(2026, 4, 22) + spend_log.metadata = { + "guardrail_information": [ + {"guardrail_name": "db-1", "guardrail_status": "not_run", "duration": 0.0}, + ] + } + prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log]) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_logs( + guardrail_id="db-1", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + assert [log.action for log in resp.logs] == ["not_run"] + + +@pytest.mark.asyncio +async def test_logs_action_passed_filter_excludes_not_run_entries(): + """LIT-6314: filtering the drill-down for passes must not return unscanned requests.""" + index_row = MagicMock() + index_row.request_id = "req-nr" + index_row.guardrail_id = "db-1" + index_row.start_time = datetime(2026, 4, 22) + spend_log = MagicMock() + spend_log.request_id = "req-nr" + spend_log.model = "gpt-4o-mini" + spend_log.startTime = datetime(2026, 4, 22) + spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "not_run"}]} + prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log]) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_logs( + guardrail_id="db-1", + policy_id=None, + page=1, + page_size=50, + action="passed", + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + assert resp.logs == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index dcbe515d5de..8382a5ada96 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -2,10 +2,8 @@ Unit tests for compliance check endpoints (EU AI Act and GDPR). """ - import pytest - from litellm.proxy.compliance_checks import ComplianceChecker from litellm.types.proxy.compliance_endpoints import ComplianceCheckRequest @@ -591,3 +589,37 @@ class TestModeMatching: continue if matched: assert mode in _guaranteed_modes(g_mode), (g_mode, mode) + + +class TestNotRunGuardrails: + """LIT-6314 logs a not_run entry for a guardrail that message scoping left nothing to scan.""" + + def test_not_run_alone_never_evidences_compliance(self): + data = ComplianceCheckRequest( + request_id="req-601", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T00:00:00Z", + guardrail_information=[ + {"guardrail_name": "pii_detection", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, + ], + ) + results = {c.check_name: c.passed for c in ComplianceChecker(data).check_eu_ai_act()} + assert results["Guardrails applied"] is False + assert results["Content screened before LLM"] is False + assert results["Audit record complete"] is False + + def test_not_run_sibling_does_not_fail_a_passing_request(self): + data = ComplianceCheckRequest( + request_id="req-602", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T00:00:00Z", + pii_detected=True, + guardrail_information=[ + {"guardrail_name": "pii_detection", "guardrail_status": "success", "guardrail_mode": "pre_call"}, + {"guardrail_name": "system_only", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, + ], + ) + results = {c.check_name: c.passed for c in ComplianceChecker(data).check_gdpr()} + assert results["Sensitive data protected"] is True diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx index 8d073feae82..8b4684e87b0 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx @@ -1,4 +1,4 @@ -import { CircleCheck, ChevronDown, TriangleAlert, X } from "lucide-react"; +import { CircleCheck, ChevronDown, MinusCircle, TriangleAlert, X } from "lucide-react"; import { useQuery } from "@tanstack/react-query"; import moment from "moment"; import React, { useState } from "react"; @@ -10,9 +10,16 @@ import type { LogEntry as ViewLogsLogEntry } from "@/components/view_logs/column import type { LogEntry } from "./mockData"; const actionConfig: Record< - "blocked" | "passed" | "flagged", + "blocked" | "passed" | "flagged" | "not_run", { icon: React.ElementType; color: string; bg: string; border: string; label: string } > = { + not_run: { + icon: MinusCircle, + color: "text-muted-foreground", + bg: "bg-muted", + border: "border-border", + label: "Not run", + }, blocked: { icon: X, color: "text-destructive", diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts index 7d99ebe7c44..717f0f459e4 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts @@ -43,7 +43,7 @@ export interface LogEntry { input_snippet?: string; output_snippet?: string; score?: number; - action: "blocked" | "passed" | "flagged"; + action: "blocked" | "passed" | "flagged" | "not_run"; model?: string; reason?: string; latency_ms?: number; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index a679dc49427..961aad7dc9e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -1,7 +1,7 @@ import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; -import { LogDetailContent } from "./LogDetailContent"; +import { GuardrailJumpLink, LogDetailContent } from "./LogDetailContent"; import type { LogEntry } from "../columns"; vi.mock("../GuardrailViewer/GuardrailViewer", () => ({ @@ -489,3 +489,30 @@ describe("LogDetailContent", () => { expect(within(descriptions).getByText("-")).toBeInTheDocument(); }); }); + +describe("GuardrailJumpLink", () => { + it("does not render a not_run entry as a failure", () => { + render( + , + ); + expect(screen.getByText(/2 guardrails/)).toHaveTextContent("✓"); + expect(screen.getByText(/2 guardrails/)).not.toHaveTextContent("✗"); + }); + + it("still renders a real failure as failed", () => { + render( + , + ); + expect(screen.getByText(/2 guardrails/)).toHaveTextContent("✗"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 4c5c7b7b43f..105ae4add45 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -636,7 +636,9 @@ function RequestResponseSection({ } export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[] }) { - const allPassed = guardrailEntries.every((e) => { + // a not_run entry never evaluated the request, so it neither passes nor fails the banner + const evaluated = guardrailEntries.filter((e) => (e?.guardrail_status || e?.status) !== "not_run"); + const allPassed = evaluated.every((e) => { const status = e?.guardrail_status || e?.status; return status === "pass" || status === "passed" || status === "success"; }); From d5cf5640b0d2f5451af553024c135e25a9b05c8d Mon Sep 17 00:00:00 2001 From: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:30:29 +0800 Subject: [PATCH 008/297] fix(responses): preserve provider affinity Co-authored-by: Bytechoreographer --- .../encrypted_content_affinity_check.py | 30 +++- .../test_encrypted_content_affinity_check.py | 136 +++++++++++++++++- 2 files changed, 162 insertions(+), 4 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b623e31ce06..eba2f36b675 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,6 +37,7 @@ Safe to enable globally: """ import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx @@ -48,6 +49,7 @@ from litellm.exceptions import ( ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues @@ -161,11 +163,13 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _encryption_boundary_key( litellm_params: object, - ) -> tuple | None: + ) -> tuple[object, object] | None: """ ``(api_base, api_key)`` pair identifying an Azure resource. Two deployments sharing both are interchangeable for ``encrypted_content`` follow-ups; Azure rejects content produced by any other resource. + Missing values are resolved from ``litellm_credential_name`` without + modifying the deployment, and explicit deployment values take precedence. Accepts any object exposing dict-style ``.get(key, default)``: plain dicts (the common case in ``healthy_deployments``) as well as @@ -180,9 +184,29 @@ class EncryptedContentAffinityCheck(CustomLogger): return None api_base: Final = getter("api_base") api_key: Final = getter("api_key") - if not api_base or not api_key: + credential_name: Final = getter("litellm_credential_name") + credential_values: Final[Mapping[str, object] | None] = ( + CredentialAccessor.get_credential_values(credential_name) + if isinstance(credential_name, str) and credential_name and (api_base is None or api_key is None) + else None + ) + effective_api_base: Final = ( + api_base + if api_base is not None + else credential_values.get("api_base") + if credential_values is not None + else None + ) + effective_api_key: Final = ( + api_key + if api_key is not None + else credential_values.get("api_key") + if credential_values is not None + else None + ) + if not effective_api_base or not effective_api_key: return None - return (api_base, api_key) + return (effective_api_base, effective_api_key) def _find_deployments_on_same_encryption_boundary( self, diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index dac991a41c4..46e67983c33 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -21,8 +21,8 @@ from unittest.mock import AsyncMock, patch import pytest - import litellm +from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -1148,6 +1148,140 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): ) +def test_boundary_key_resolves_missing_values_from_named_credential(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], + ) + ): + boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) + + assert boundary == ("https://account-a.example.com", "credential-key-a") + + +def test_boundary_key_prefers_explicit_values_over_named_credential(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], + ) + ): + boundary = EncryptedContentAffinityCheck._encryption_boundary_key( + { + "api_base": "https://deployment.example.com", + "litellm_credential_name": "account-a", + } + ) + + assert boundary == ("https://deployment.example.com", "credential-key-a") + + +def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-a-peer", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-b", + credential_values={ + "api_base": "https://account-b.example.com", + "api_key": "credential-key-b", + }, + credential_info={}, + ), + ], + ) + ): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.3-codex", + "litellm_params": { + "model": "azure/gpt-5.3-codex", + "litellm_credential_name": "account-a", + }, + "model_info": {"id": "origin"}, + } + ], + num_retries=0, + ) + check = EncryptedContentAffinityCheck(router=router) + healthy_deployments = [ + { + "model_info": {"id": "peer-same-boundary"}, + "litellm_params": { + "model": "azure/gpt-5.4", + "litellm_credential_name": "account-a-peer", + }, + }, + { + "model_info": {"id": "peer-different-boundary"}, + "litellm_params": { + "model": "azure/gpt-5.4", + "litellm_credential_name": "account-b", + }, + }, + ] + + matches, originating = check._find_deployments_on_same_encryption_boundary( + healthy_deployments=healthy_deployments, + model_id="origin", + ) + + assert originating is not None + assert [deployment["model_info"]["id"] for deployment in matches] == ["peer-same-boundary"] + + def test_boundary_key_rejects_non_dict_like_inputs(): """ Inputs that don't expose ``.get()`` (None, lists, strings, ints) -> None. From 2cd28b97f1793ef6032526a0119cb892dc3e9b66 Mon Sep 17 00:00:00 2001 From: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:02:31 +0800 Subject: [PATCH 009/297] fix(responses): align credential boundary resolution --- .../encrypted_content_affinity_check.py | 26 +++++-------- .../test_encrypted_content_affinity_check.py | 37 ++++++++++++++++++- 2 files changed, 45 insertions(+), 18 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index eba2f36b675..db3c865e183 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -165,11 +165,9 @@ class EncryptedContentAffinityCheck(CustomLogger): litellm_params: object, ) -> tuple[object, object] | None: """ - ``(api_base, api_key)`` pair identifying an Azure resource. Two - deployments sharing both are interchangeable for ``encrypted_content`` - follow-ups; Azure rejects content produced by any other resource. - Missing values are resolved from ``litellm_credential_name`` without - modifying the deployment, and explicit deployment values take precedence. + ``(api_base, api_key)`` identifies an upstream encryption boundary. + The values are resolved from the deployment and its named credential + without modifying the deployment. Accepts any object exposing dict-style ``.get(key, default)``: plain dicts (the common case in ``healthy_deployments``) as well as @@ -187,22 +185,18 @@ class EncryptedContentAffinityCheck(CustomLogger): credential_name: Final = getter("litellm_credential_name") credential_values: Final[Mapping[str, object] | None] = ( CredentialAccessor.get_credential_values(credential_name) - if isinstance(credential_name, str) and credential_name and (api_base is None or api_key is None) + if isinstance(credential_name, str) and credential_name else None ) effective_api_base: Final = ( - api_base - if api_base is not None - else credential_values.get("api_base") - if credential_values is not None - else None + credential_values.get("api_base") + if credential_values is not None and "api_base" in credential_values + else api_base ) effective_api_key: Final = ( - api_key - if api_key is not None - else credential_values.get("api_key") - if credential_values is not None - else None + credential_values.get("api_key") + if credential_values is not None and "api_key" in credential_values + else api_key ) if not effective_api_base or not effective_api_key: return None diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 46e67983c33..7b7a4969d41 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1174,7 +1174,7 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): assert boundary == ("https://account-a.example.com", "credential-key-a") -def test_boundary_key_prefers_explicit_values_over_named_credential(): +def test_boundary_key_matches_named_credential_precedence(): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1198,11 +1198,44 @@ def test_boundary_key_prefers_explicit_values_over_named_credential(): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { "api_base": "https://deployment.example.com", + "api_key": "deployment-key", "litellm_credential_name": "account-a", } ) - assert boundary == ("https://deployment.example.com", "credential-key-a") + assert boundary == ("https://credential.example.com", "credential-key-a") + + +def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], + ) + ): + boundary = EncryptedContentAffinityCheck._encryption_boundary_key( + { + "api_base": "", + "api_key": "", + "litellm_credential_name": "account-a", + } + ) + + assert boundary == ("https://credential.example.com", "credential-key-a") def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): From 990dea27d5c87c7c48dbc286c2efa3c6a610cf54 Mon Sep 17 00:00:00 2001 From: Rad Wadud <104943953+rad-p44@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:52:41 -0500 Subject: [PATCH 010/297] fix(headroom): protect cache_control-marked rows anywhere in history get_protected_indices() only protected system rows, the last user row, and the last assistant row. A message carrying its own Anthropic cache_control breakpoint further back in history (e.g. a large cached tool result from a few turns ago) was not protected, so the Headroom guardrail would send it to /v1/compress and rewrite it. The row came back byte-different but kept its cache_control marker, so the provider's prompt cache treated the next request as a miss on that prefix: a cache read silently became a cache write. This reproduces the production cache-hit-rate collapse reported in #39519 (~65-70% down to ~40-50% within 48h of enabling the guardrail). get_protected_indices() now also protects any message whose content -- directly on the message, or on any part of a list-of-parts content -- carries a cache_control marker, regardless of its position in history. Both compress() and the Headroom guardrail already share this function as their compression-eligibility policy, so both get the fix. Adds test coverage for cache_control on the message dict itself, on a content part, mid-history, and de-duplicated against already-protected indices. Updates the Headroom guardrail's PARTS_MESSAGES fixture, which previously relied on this exact gap for its all-text merge/flatten test coverage, to use a separate un-marked row (the cache_control-marked-row merge behavior is covered directly by compresr's own test, since a marked row no longer reaches that merge path through Headroom). Fixes #39519 --- litellm/compression/compress.py | 35 +++++++++- .../test_litellm/compression/test_compress.py | 60 ++++++++++++++++ .../guardrail_hooks/test_headroom.py | 68 ++++++++++++++++--- 3 files changed, 151 insertions(+), 12 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index c646baf9d9e..62b05a4938f 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -205,21 +205,54 @@ def _extract_anthropic_tool_exchange_spans( return spans, None +def _message_has_cache_control(message: Mapping[str, object]) -> bool: + """True if ``message`` carries an Anthropic ``cache_control`` breakpoint. + + A breakpoint can sit directly on the message dict, or on any part of a + list-of-parts ``content`` (the shape Anthropic's own messages use). Either + placement pins the provider's KV-cache prefix to this row's exact bytes, so + either placement must protect the row the same way. + """ + if message.get("cache_control") is not None: + return True + content: Final = message.get("content") + if isinstance(content, list): + return any(isinstance(part, Mapping) and part.get("cache_control") is not None for part in content) + return False + + def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]: """ Return indices of messages that must never be compressed: - All system messages - The last user message - The last assistant message + - Any message carrying an Anthropic cache_control breakpoint The last user message is what the model is being asked to act on right now, so compressing it replaces the live instruction with a marker. Compression guardrails share this policy; see the Headroom guardrail. + + A cache_control breakpoint pins the provider's prompt-cache prefix to that + row's exact bytes. Rewriting the row (even leaving the marker in place) + changes those bytes, so the next request misses the cache it thinks it is + reusing and silently pays a cache write instead of a cache read. This is + not limited to the last user/assistant row: a marker several turns back + (e.g. on a large cached tool result) needs the same protection. """ system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system") last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:] last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] - return system_indices + last_user + last_assistant + cache_control_indices: Final = tuple( + index for index, msg in enumerate(messages) if _message_has_cache_control(msg) + ) + seen: Final[set[int]] = set() + ordered: Final[list[int]] = [] + for index in system_indices + last_user + last_assistant + cache_control_indices: + if index not in seen: + seen.add(index) + ordered.append(index) + return tuple(ordered) def _combine_scores( diff --git a/tests/test_litellm/compression/test_compress.py b/tests/test_litellm/compression/test_compress.py index 6827c37dfd5..f9877ea2bc4 100644 --- a/tests/test_litellm/compression/test_compress.py +++ b/tests/test_litellm/compression/test_compress.py @@ -53,3 +53,63 @@ def test_every_system_row_is_protected(): def test_no_user_or_assistant_rows(): assert sorted(get_protected_indices([{"role": "system", "content": "sys"}])) == [0] assert get_protected_indices([]) == () + + +def test_mid_history_cache_control_part_is_protected(): + # A large cached tool result from a few turns back, not the last user or + # last assistant row -- exactly the row a provider prompt-cache pins to + # exact bytes. Rewriting it (even leaving the marker on) changes those + # bytes and turns the next request's cache read into a cache write. + messages = [ + {"role": "user", "content": "old question"}, + {"role": "assistant", "content": "old answer"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "a large cached tool result"}, + ], + }, + {"role": "assistant", "content": "ack"}, + {"role": "user", "content": "live instruction"}, + ] + messages[2]["content"][0]["cache_control"] = {"type": "ephemeral"} + + # index 3 = last assistant, index 4 = last user (both protected by role + # regardless), index 2 = the cache_control-marked row itself. + assert sorted(get_protected_indices(messages)) == [2, 3, 4] + + +def test_cache_control_directly_on_message_is_protected(): + messages = [ + {"role": "user", "content": "old question", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "old answer"}, + {"role": "user", "content": "live instruction"}, + ] + + assert sorted(get_protected_indices(messages)) == [0, 1, 2] + + +def test_cache_control_protection_does_not_duplicate_already_protected_rows(): + # The last user row is already protected by role; marking it too must not + # produce a duplicate index. + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "live", "cache_control": {"type": "ephemeral"}}, + ] + + protected = get_protected_indices(messages) + + assert sorted(protected) == [0, 1] + assert len(protected) == len(set(protected)) + + +def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control(): + # Defensive: a plain string content, or a list of non-dict items, must not + # raise or be misread as carrying a breakpoint. + messages = [ + {"role": "assistant", "content": "plain string content"}, + {"role": "user", "content": ["not", "a", "dict", "list"]}, + {"role": "user", "content": "live instruction"}, + ] + + assert sorted(get_protected_indices(messages)) == [0, 2] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index d8eeb8d2b8a..5cd42bd3f83 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1795,14 +1795,18 @@ PARTS_MESSAGES = [ ], }, { + # No cache_control here on purpose: this row exercises the general + # multi-part flatten/merge mechanics (shared with compresr). A row + # carrying its own cache_control is a different, dedicated case -- + # see test_mid_history_cache_control_row_is_never_sent_for_compression + # (#39519): get_protected_indices withholds it from /v1/compress + # entirely rather than letting it be rewritten and re-merged, because + # rewriting the bytes under a live breakpoint busts the cache the + # marker is supposed to preserve. "role": "user", "content": [ - {"type": "text", "text": "Earlier turn.", "cache_control": {"type": "ephemeral"}}, - { - "type": "text", - "text": "Second block. " + "B" * 5000, - "cache_control": {"type": "ephemeral", "ttl": "1h"}, - }, + {"type": "text", "text": "Earlier turn."}, + {"type": "text", "text": "Second block. " + "B" * 5000}, ], }, { @@ -1891,14 +1895,17 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( messages = result["structured_messages"] history_content = messages[1]["content"] - # Rewritten all-text row collapses to one part carrying the LAST declared - # breakpoint: an Anthropic breakpoint caches the prefix ending at its - # part, so after the merge the last one (and its TTL) still describes the - # row. + # Rewritten all-text row collapses to one part carrying the rewritten text. + # This fixture row carries no cache_control (see PARTS_MESSAGES): the + # last-declared-breakpoint-survives-the-merge behavior is a property of + # merge_rewritten_text_parts and is covered directly by compresr's + # test_all_text_row_merges_and_keeps_last_cache_control, since a + # cache_control-marked row never reaches this merge path through Headroom + # at all -- get_protected_indices withholds it before compression runs + # (see test_mid_history_cache_control_row_is_never_sent_for_compression). assert isinstance(history_content, list) assert len(history_content) == 1 assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac" - assert history_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} # Mixed row passes through byte-identical. assert messages[2]["content"] == PARTS_MESSAGES[2]["content"] # The service-declared hash still drives retrieve-tool injection on a restored row. @@ -2523,6 +2530,45 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert messages[3] == compressed_history[1] +# --------------------------------------------------------------------------- +# #39519: a mid-history row carrying its own Anthropic cache_control marker +# (e.g. a large tool result the client already cached several turns back) was +# still sent to /v1/compress and rewritten. It came back byte-different but +# kept its marker, so the next request's cache read silently became a cache +# write. get_protected_indices() now protects any cache_control-marked row, +# not just system/last-user/last-assistant, so it must never reach the wire. +# --------------------------------------------------------------------------- + +CACHE_MARKED_HISTORY_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "old question " + "Q" * 5000}, + { + "role": "assistant", + "content": "Reading the file now.", + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + { + "role": "tool", + "tool_call_id": "old_1", + "content": [{"type": "text", "text": "large cached file body " + "F" * 5000}], + "cache_control": {"type": "ephemeral"}, + }, + {"role": "assistant", "content": "Summarized the file for you."}, + {"role": "user", "content": "live instruction"}, +] + + +@pytest.mark.asyncio +async def test_mid_history_cache_control_row_is_never_sent_for_compression(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, CACHE_MARKED_HISTORY_MESSAGES) + + cached_row = CACHE_MARKED_HISTORY_MESSAGES[3] + assert cached_row not in wire + assert not any(row.get("tool_call_id") == "old_1" for row in wire) + # Byte-identical, marker intact -- the next request's cache read survives. + assert result["structured_messages"][3] == cached_row + + # --------------------------------------------------------------------------- # #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP # gateway) executes headroom_retrieve and echoes the recovered original content From e2b41286d351c6eda893c7e335d4a03896791f5e Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 00:48:27 -0400 Subject: [PATCH 011/297] fix(vertex-live): bill every modality on the /vertex_ai/live passthrough The Live passthrough builds Usage from the TEXT-modality counts alone, so audio, image and video tokens never reach the cost calculator and bill as nothing. A one-turn audio session reported 13 text and 127 audio input tokens and billed the 13; a camera session reported 1043 prompt tokens and billed 11. Reporting the full per-modality breakdown fixes it, because the shared Gemini input and output cost path already prices audio, image and video from prompt_tokens_details and completion_tokens_details. On the native-audio entry that is a 6x difference per token in both directions, which is the whole gap. Aggregation across turns is unchanged. Google charges per turn for every token in the Live session context window, current turn plus all accumulated tokens from previous turns, so the existing summing is what Vertex bills and it stays as it is. That is worth stating because the cumulative promptTokensDetails looks like a restatement of one running total, and treating it that way would under-bill a multi-turn session. See the LiveAPI context-window note on https://cloud.google.com/vertex-ai/generative-ai/pricing. Live can also name the modality carrying the rest of a turn and omit its tokenCount. Reading that absent key as zero left the tokens inside candidatesTokenCount but outside the breakdown, so real speech was charged at the text output rate. A lone unpriced entry now takes whatever the turn's declared count leaves over. Two or more cannot be told apart, so they are still left to the calculator's text remainder. Server-side toolUsePromptTokenCount is now reported in prompt_tokens_details. It is deliberately kept out of prompt_tokens: no Gemini route prices tool-use tokens, and adding them there instead suppresses the cache-overlap correction and raises the bill for no extra work. Removes _calculate_live_api_cost, whose result never reached the bill. It set kwargs["response_cost"], which the standard logging path recomputes from the ModelResponse, and on a measured audio session it returned $0.000487 against a $0.0000425 row. Now that the modality counts reach the standard calculator, keeping a second hand-rolled pricing path would only ever double-charge. The rewrite of the aggregator is arithmetically identical to what it replaced. It sums the same three counts and the same per-modality details, still takes the remaining fields from the first turn, and drops nine LIT010, one C901 and 42 basedpyright findings in the process. --- ...tex_ai_live_passthrough_logging_handler.py | 335 +++++++----------- .../test_vertex_ai_live_passthrough.py | 326 ++++++++++++----- 2 files changed, 375 insertions(+), 286 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index e26f5f57532..e224f707b02 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -5,7 +5,10 @@ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough e Supports different modalities: text, audio, video, and web search. """ +from collections.abc import Mapping, Sequence from datetime import datetime +from itertools import chain +from types import MappingProxyType from typing import Any, Final from litellm._logging import verbose_proxy_logger @@ -15,8 +18,23 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( PassThroughEndpointLoggingTypedDict, ) -from litellm.types.utils import LlmProviders, ModelResponse, Usage -from litellm.utils import get_model_info +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + LlmProviders, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +_AGGREGATED_FIELDS: Final = frozenset( + { + "promptTokenCount", + "candidatesTokenCount", + "totalTokenCount", + "promptTokensDetails", + "candidatesTokensDetails", + } +) class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @@ -48,6 +66,56 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """Return the LLM provider name.""" return LlmProviders.VERTEX_AI + @staticmethod + def _resolve_detail_counts( + details: Sequence[Mapping[str, Any]], + declared_total: object, + ) -> tuple[tuple[str, int], ...]: + """ + Pair each of one turn's ``*TokensDetails`` entries with its token count. + + Live sometimes names the modality that carries the rest of a turn without a + ``tokenCount``, and reading the absent key as zero drops those tokens from the + breakdown, so real audio ends up priced as text. A lone unpriced entry therefore takes + whatever the turn's declared count leaves over. Two or more cannot be told apart, so + they are left out and the cost calculator charges the remainder as text. + """ + priced: Final = tuple( + (str(detail.get("modality", "TEXT")), count) + for detail in details + if isinstance(count := detail.get("tokenCount"), int) + ) + unpriced: Final = tuple( + str(detail.get("modality", "TEXT")) for detail in details if not isinstance(detail.get("tokenCount"), int) + ) + if len(unpriced) != 1 or not isinstance(declared_total, int): + return priced + residual: Final = declared_total - sum(count for _, count in priced) + return priced if residual <= 0 else (*priced, (unpriced[0], residual)) + + @staticmethod + def _sum_by_modality(counts: Sequence[tuple[str, int]]) -> Mapping[str, int]: + """Total the (modality, tokenCount) pairs of one or more turns per modality.""" + return MappingProxyType({modality: sum(c for m, c in counts if m == modality) for modality, _ in counts}) + + @staticmethod + def _merged_modality_totals( + snapshots: Sequence[Mapping[str, Any]], + count_key: str, + details_key: str, + ) -> Mapping[str, int]: + """Total every turn's per-modality counts, so the breakdown adds up the way the totals do.""" + return VertexAILivePassthroughLoggingHandler._sum_by_modality( + tuple( + chain.from_iterable( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + snapshot.get(details_key) or [], snapshot.get(count_key) + ) + for snapshot in snapshots + ) + ) + ) + @staticmethod def _extract_usage_metadata_from_websocket_messages( websocket_messages: list[dict], @@ -55,175 +123,45 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Extract and aggregate usage metadata from a list of WebSocket messages. + Live emits one ``usageMetadata`` per turn and Google charges per turn for every token in + the session context window, which is the current turn's tokens plus all accumulated + tokens from previous turns, so the turns add up rather than restating each other. See + the Live API note under https://cloud.google.com/vertex-ai/generative-ai/pricing. + Args: websocket_messages: List of WebSocket messages from the Live API Returns: Dictionary containing aggregated usage metadata, or None if not found """ - all_usage_metadata: Final = [] + snapshots: Final = tuple( + message["usageMetadata"] + for message in websocket_messages + if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict) + ) - # Collect all usage metadata messages - for message in websocket_messages: - if isinstance(message, dict) and "usageMetadata" in message: - all_usage_metadata.append(message["usageMetadata"]) - - if not all_usage_metadata: + if not snapshots: return None - # If only one usage metadata, return it as-is - if len(all_usage_metadata) == 1: - return all_usage_metadata[0] - - # Aggregate multiple usage metadata messages - aggregated: Final[dict[str, Any]] = { - "promptTokenCount": 0, - "candidatesTokenCount": 0, - "totalTokenCount": 0, - "promptTokensDetails": [], - "candidatesTokensDetails": [], + prompt_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals( + snapshots, "promptTokenCount", "promptTokensDetails" + ) + candidate_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals( + snapshots, "candidatesTokenCount", "candidatesTokensDetails" + ) + return { + **{key: value for key, value in snapshots[0].items() if key not in _AGGREGATED_FIELDS}, + "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), + "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), + "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "promptTokensDetails": [ + {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 + ], + "candidatesTokensDetails": [ + {"modality": modality, "tokenCount": count} for modality, count in candidate_totals.items() if count > 0 + ], } - # Aggregate token counts - for usage in all_usage_metadata: - aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0) - aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0) - aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0) - - # Aggregate token details by modality - modality_totals: Final = {} - - for usage in all_usage_metadata: - # Process prompt tokens details - for detail in usage.get("promptTokensDetails", []): - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality not in modality_totals: - modality_totals[modality] = {"prompt": 0, "candidate": 0} - modality_totals[modality]["prompt"] += token_count - - # Process candidate tokens details - for detail in usage.get("candidatesTokensDetails", []): - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality not in modality_totals: - modality_totals[modality] = {"prompt": 0, "candidate": 0} - modality_totals[modality]["candidate"] += token_count - - # Convert aggregated modality totals back to details format - for modality, totals in modality_totals.items(): - if totals["prompt"] > 0: - aggregated["promptTokensDetails"].append({"modality": modality, "tokenCount": totals["prompt"]}) - if totals["candidate"] > 0: - aggregated["candidatesTokensDetails"].append({"modality": modality, "tokenCount": totals["candidate"]}) - - # Add any additional fields from the first usage metadata - first_usage: Final = all_usage_metadata[0] - for key, value in first_usage.items(): - if key not in aggregated: - aggregated[key] = value - - return aggregated - - @staticmethod - def _calculate_live_api_cost( - model: str, - usage_metadata: dict, - custom_llm_provider: str = "vertex_ai", - ) -> float: - """ - Calculate cost for Vertex AI Live API based on usage metadata. - - Args: - model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09") - usage_metadata: Usage metadata from the Live API response - custom_llm_provider: The LLM provider (default: "vertex_ai") - - Returns: - Total cost in USD - """ - try: - # Get model pricing information - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - - verbose_proxy_logger.debug("Vertex AI Live API model info for '%s': %s", model, model_info) - - # Check if pricing info is available - if not model_info or not model_info.get("input_cost_per_token"): - verbose_proxy_logger.error("No pricing info found for %s in local model pricing database", model) - return 0.0 - - total_cost = 0.0 - - # Extract token counts from usage metadata - prompt_token_count: Final = usage_metadata.get("promptTokenCount", 0) - candidates_token_count: Final = usage_metadata.get("candidatesTokenCount", 0) - - # Calculate base text token costs - input_cost_per_token: Final = model_info.get("input_cost_per_token", 0.0) - output_cost_per_token: Final = model_info.get("output_cost_per_token", 0.0) - - total_cost += prompt_token_count * input_cost_per_token - total_cost += candidates_token_count * output_cost_per_token - - # Handle modality-specific costs if present - prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", []) - - # Process prompt tokens by modality - for detail in prompt_tokens_details: - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality == "AUDIO": - audio_cost_per_token = model_info.get("input_cost_per_audio_token", 0.0) - total_cost += token_count * audio_cost_per_token - elif modality == "VIDEO": - # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get("input_cost_per_video_per_second", 0.0) - total_cost += token_count * video_cost_per_token - # TEXT tokens are already handled above - - # Process candidate tokens by modality - for detail in candidates_tokens_details: - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality == "AUDIO": - audio_cost_per_token = model_info.get("output_cost_per_audio_token", 0.0) - total_cost += token_count * audio_cost_per_token - elif modality == "VIDEO": - # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get("output_cost_per_video_per_second", 0.0) - total_cost += token_count * video_cost_per_token - # TEXT tokens are already handled above - - # Handle web search costs if present - tool_use_prompt_token_count: Final = usage_metadata.get("toolUsePromptTokenCount", 0) - if tool_use_prompt_token_count > 0: - # Web search typically has a fixed cost per request - web_search_cost: Final = model_info.get("web_search_cost_per_request", 0.0) - if isinstance(web_search_cost, (int, float)) and web_search_cost > 0: - total_cost += web_search_cost - else: - # Fallback to token-based pricing for tool use - total_cost += tool_use_prompt_token_count * input_cost_per_token - - verbose_proxy_logger.debug( - f"Vertex AI Live API cost calculation - Model: {model}, " - f"Prompt tokens: {prompt_token_count}, " - f"Candidate tokens: {candidates_token_count}, " - f"Total cost: ${total_cost:.6f}" - ) - - return total_cost - - except Exception as e: - verbose_proxy_logger.error("Error calculating Vertex AI Live API cost: %s", e) - return 0.0 - @staticmethod def _create_usage_object_from_metadata( usage_metadata: dict, @@ -239,38 +177,37 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Returns: LiteLLM Usage object """ - prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) - completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) - total_tokens: Final = usage_metadata.get("totalTokenCount", 0) + prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount") + ) + ) + candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount") + ) + ) - # Create modality-specific token details if available - prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", []) - - # Extract text tokens from details - text_prompt_tokens = 0 - text_completion_tokens = 0 - - for detail in prompt_tokens_details: - if detail.get("modality") == "TEXT": - text_prompt_tokens = detail.get("tokenCount", 0) - break - - for detail in candidates_tokens_details: - if detail.get("modality") == "TEXT": - text_completion_tokens = detail.get("tokenCount", 0) - break - - # If no text tokens found in details, use total counts - if text_prompt_tokens == 0: - text_prompt_tokens = prompt_tokens - if text_completion_tokens == 0: - text_completion_tokens = completion_tokens + prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) + completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) return Usage( - prompt_tokens=text_prompt_tokens, - completion_tokens=text_completion_tokens, - total_tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=prompt_by_modality.get("TEXT"), + audio_tokens=prompt_by_modality.get("AUDIO"), + image_tokens=prompt_by_modality.get("IMAGE"), + video_tokens=prompt_by_modality.get("VIDEO"), + tool_use_tokens=usage_metadata.get("toolUsePromptTokenCount") or None, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=candidates_by_modality.get("TEXT"), + audio_tokens=candidates_by_modality.get("AUDIO"), + image_tokens=candidates_by_modality.get("IMAGE"), + video_tokens=candidates_by_modality.get("VIDEO"), + ), ) def vertex_ai_live_passthrough_handler( @@ -316,13 +253,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "kwargs": kwargs, } - # Calculate cost using Live API specific pricing - response_cost: Final = self._calculate_live_api_cost( - model=model, - usage_metadata=usage_metadata, - custom_llm_provider=custom_llm_provider, - ) - # Create Usage object for standard LiteLLM logging usage: Final = self._create_usage_object_from_metadata( usage_metadata=usage_metadata, @@ -339,8 +269,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): choices=[], ) - # Update kwargs with cost information - kwargs["response_cost"] = response_cost kwargs["model"] = model kwargs["custom_llm_provider"] = custom_llm_provider @@ -350,10 +278,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$") safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( - f"Vertex AI Live API passthrough cost tracking - " - f"Model: {safe_model}, Cost: ${response_cost:.6f}, " - f"Prompt tokens: {usage.prompt_tokens}, " - f"Completion tokens: {usage.completion_tokens}" + "Vertex AI Live API passthrough cost tracking - Model: %s, " + "Prompt tokens: %s %s, Completion tokens: %s %s", + safe_model, + usage.prompt_tokens, + usage.prompt_tokens_details, + usage.completion_tokens, + usage.completion_tokens_details, ) return { diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index e2eb6d0b68b..3b6a548b219 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -201,88 +201,247 @@ class TestVertexAILivePassthroughLoggingHandler: assert text_prompt["tokenCount"] == 10 assert audio_prompt["tokenCount"] == 10 - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_calculate_cost_basic(self, mock_get_model_info, handler): - """Test basic cost calculation""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - } + def test_usage_carries_every_modality(self, handler): + """Regression: the Usage object reported only TEXT, so audio and image billed as nothing. + prompt_tokens must be the full count and the details must name each modality, + because the cost calculator prices audio and image from *_tokens_details. + """ usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - } - - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) - - # The cost calculation may include additional factors, so we check it's reasonable - expected_min_cost = (100 * 0.000001) + (50 * 0.000002) - assert cost >= expected_min_cost - assert cost > 0 - - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_calculate_cost_with_audio(self, mock_get_model_info, handler): - """Test cost calculation with audio tokens""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - "input_cost_per_audio_token": 0.0001, - "output_cost_per_audio_token": 0.0002, - } - - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, + "promptTokenCount": 1300, + "candidatesTokenCount": 124, + "totalTokenCount": 1424, "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 80}, - {"modality": "AUDIO", "tokenCount": 20}, + {"modality": "TEXT", "tokenCount": 13}, + {"modality": "AUDIO", "tokenCount": 127}, + {"modality": "IMAGE", "tokenCount": 1160}, ], "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 30}, - {"modality": "AUDIO", "tokenCount": 20}, + {"modality": "TEXT", "tokenCount": 29}, + {"modality": "AUDIO", "tokenCount": 95}, ], } - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + usage = handler._create_usage_object_from_metadata( + usage_metadata=usage_metadata, model="gemini-live-2.5-flash" + ) - # Should include both text and audio costs - assert cost > 0 - assert cost > (100 * 0.000001) + ( - 50 * 0.000002 - ) # Should be higher due to audio + assert usage.prompt_tokens == 1300, "the full prompt count must survive, not just its text share" + assert usage.completion_tokens == 124 + assert usage.prompt_tokens_details.text_tokens == 13 + assert usage.prompt_tokens_details.audio_tokens == 127 + assert usage.prompt_tokens_details.image_tokens == 1160 + assert usage.completion_tokens_details.text_tokens == 29 + assert usage.completion_tokens_details.audio_tokens == 95 - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" + def test_usage_sums_repeated_modality_entries(self, handler): + """A modality can appear more than once across aggregated turns; sum, don't overwrite.""" + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 40, + "candidatesTokenCount": 0, + "promptTokensDetails": [ + {"modality": "IMAGE", "tokenCount": 10}, + {"modality": "IMAGE", "tokenCount": 25}, + {"modality": "TEXT", "tokenCount": 5}, + ], + }, + model="gemini-live-2.5-flash", + ) + assert usage.prompt_tokens_details.image_tokens == 35 + assert usage.prompt_tokens_details.text_tokens == 5 + + NATIVE_AUDIO_MODEL = "gemini-live-2.5-flash-preview-native-audio-09-2025" + + # A four-turn native-audio session. Google charges per turn for the whole session context + # window, so the prompt side repeats the accumulated audio while the candidates side reports + # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is + # the shape Live really emits at the end of a spoken answer. + AUDIO_SESSION = ( + {"prompt": (14, 122), "candidates": (8, 20)}, + {"prompt": (21, 182), "candidates": (5, 50)}, + {"prompt": (24, 203), "candidates": (13, 27)}, + {"prompt": (24, 203), "candidates": (0, 3), "candidate_audio_token_count_missing": True}, ) - def test_calculate_cost_with_web_search(self, mock_get_model_info, handler): - """Test cost calculation with web search (tool use)""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - "web_search_cost_per_request": 0.01, - } - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - "toolUsePromptTokenCount": 10, - } + @staticmethod + def _live_messages(turns): + """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits.""" + return [{"type": "session.created", "session": {"id": "s"}}] + [ + { + "type": "response.done", + "usageMetadata": { + "promptTokenCount": sum(turn["prompt"]), + "candidatesTokenCount": sum(turn["candidates"]), + "totalTokenCount": sum(turn["prompt"]) + sum(turn["candidates"]), + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": turn["prompt"][0]}, + {"modality": "AUDIO", "tokenCount": turn["prompt"][1]}, + ], + "candidatesTokensDetails": ( + [{"modality": "AUDIO"}] + if turn.get("candidate_audio_token_count_missing") + else [ + {"modality": "TEXT", "tokenCount": turn["candidates"][0]}, + {"modality": "AUDIO", "tokenCount": turn["candidates"][1]}, + ] + ), + }, + } + for turn in turns + ] - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + @staticmethod + def _session_usage(handler, mock_logging_obj, messages, model): + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=mock_logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=model, + ) + assert result["result"] is not None, "the handler must produce a usage-bearing response to bill" + return result["result"].usage - # Should include web search cost - expected_base_cost = (100 * 0.000001) + (50 * 0.000002) - # The web search cost might be handled differently, so just check it's reasonable - assert cost >= expected_base_cost - assert cost > 0 + @classmethod + def _session_cost(cls, handler, mock_logging_obj, messages, model): + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse + + usage = cls._session_usage(handler, mock_logging_obj, messages, model) + return completion_cost( + completion_response=ModelResponse( + id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[] + ), + model=f"vertex_ai/{model}", + custom_llm_provider="vertex_ai", + call_type="acompletion", + ) + + @classmethod + def _expected_session_cost(cls, turns): + from litellm.utils import get_model_info + + info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai") + return ( + sum(turn["prompt"][0] for turn in turns) * info["input_cost_per_token"] + + sum(turn["prompt"][1] for turn in turns) * info["input_cost_per_audio_token"] + + sum(turn["candidates"][0] for turn in turns) * info["output_cost_per_token"] + + sum(turn["candidates"][1] for turn in turns) * info["output_cost_per_audio_token"] + ) + + def test_every_turn_of_a_session_is_billed(self, handler, mock_logging_obj): + """Google charges per turn for the whole context window, so every turn adds to the bill. + + Billing one snapshot instead gives away all the other turns: on this session the + largest single turn is well under the session total, and its share of the audio is + priced 6x the text rate, so the gap is money rather than rounding. + """ + turns = self.AUDIO_SESSION[:3] + cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + + assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) + widest_single_turn = max(self._expected_session_cost([turn]) for turn in turns) + assert cost > widest_single_turn, "billing one snapshot drops every other turn of the session" + + def test_audio_named_without_a_token_count_bills_at_the_audio_rate(self, handler, mock_logging_obj): + """Live can name the modality carrying the rest of a turn and omit its tokenCount. + + Reading the absent key as zero left those tokens inside candidatesTokenCount but outside + the breakdown, so the calculator charged real speech at the text output rate. At this + entry's rates the last turn's 3 audio tokens are $0.0000360 rather than $0.0000060. + """ + turns = self.AUDIO_SESSION + usage = self._session_usage(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + + assert usage.completion_tokens_details.audio_tokens == 100, "the unpriced entry takes the turn's residual" + assert usage.completion_tokens_details.text_tokens == 26 + assert usage.completion_tokens == 126 + + cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) + + def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): + """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + + It is reported, not billed. Nothing in the shared Gemini input-cost path prices + tool-use tokens, and folding them into prompt_tokens here would suppress that + path's cache-overlap correction and raise the bill instead. + """ + messages = self._live_messages(self.AUDIO_SESSION[:1]) + grounded = [dict(message) for message in messages] + grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + + usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + assert usage.prompt_tokens_details.tool_use_tokens == 500 + + plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" + + @pytest.mark.parametrize( + "label,prompt_details,candidate_details", + [ + ("text only", [("TEXT", 6)], [("TEXT", 2)]), + ("audio in", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 18)]), + ("image in", [("TEXT", 10), ("IMAGE", 258)], [("TEXT", 24)]), + ("frames in", [("TEXT", 11), ("IMAGE", 1032)], [("TEXT", 26)]), + ("audio both ways", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 29), ("AUDIO", 95)]), + ], + ) + def test_live_session_bills_each_modality_at_its_own_rate(self, handler, label, prompt_details, candidate_details): + """Every payload here is a real Vertex Live session's usageMetadata. + + Before the fix these billed the text share only, from 1x (text) to 55x under. + The expected amount is derived from the entry's own rates rather than hardcoded, + so this stays correct as prices move, and it is asserted exactly, so dropping a + modality and double-charging one both fail. + """ + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse + from litellm.utils import get_model_info + + model = self.NATIVE_AUDIO_MODEL + info = get_model_info(model=model, custom_llm_provider="vertex_ai") + + text_in = info["input_cost_per_token"] + audio_in = info.get("input_cost_per_audio_token") or text_in + image_in = info.get("input_cost_per_image_token") or text_in + text_out = info["output_cost_per_token"] + audio_out = info.get("output_cost_per_audio_token") or text_out + rate_in = {"TEXT": text_in, "AUDIO": audio_in, "IMAGE": image_in} + rate_out = {"TEXT": text_out, "AUDIO": audio_out} + + expected = sum(c * rate_in[m] for m, c in prompt_details) + sum(c * rate_out[m] for m, c in candidate_details) + + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": sum(c for _, c in prompt_details), + "candidatesTokenCount": sum(c for _, c in candidate_details), + "promptTokensDetails": [{"modality": m, "tokenCount": c} for m, c in prompt_details], + "candidatesTokensDetails": [{"modality": m, "tokenCount": c} for m, c in candidate_details], + }, + model=model, + ) + + cost = completion_cost( + completion_response=ModelResponse( + id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[] + ), + model=f"vertex_ai/{model}", + custom_llm_provider="vertex_ai", + call_type="acompletion", + ) + + assert cost == pytest.approx(expected, rel=1e-9), label + + text_only = sum(c for m, c in prompt_details if m == "TEXT") * text_in + sum( + c for m, c in candidate_details if m == "TEXT" + ) * text_out + if any(m != "TEXT" for m, _ in prompt_details + candidate_details) and audio_in != text_in: + assert cost > text_only, f"{label}: non-text modalities must add cost" def test_vertex_ai_live_passthrough_handler_integration( self, handler, mock_logging_obj, sample_websocket_messages @@ -540,25 +699,24 @@ class TestVertexAILivePassthroughErrorHandling: result = handler._extract_usage_metadata_from_websocket_messages(messages) assert result is None - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_cost_calculation_with_missing_model_info(self, mock_get_model_info): - """Test cost calculation when model info is missing""" + def test_usage_without_modality_details(self): + """Older payloads carry only the totals; fall back to them rather than reporting zero.""" handler = VertexAILivePassthroughLoggingHandler() - # Mock missing model info - mock_get_model_info.return_value = {} + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + }, + model="unknown-model", + ) - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - } - - # Should not raise an exception, should return 0 or handle gracefully - cost = handler._calculate_live_api_cost("unknown-model", usage_metadata) - assert cost == 0.0 + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 50 + assert usage.total_tokens == 150 + assert usage.prompt_tokens_details.audio_tokens is None + assert usage.prompt_tokens_details.image_tokens is None def test_handler_with_none_websocket_messages(self, mock_logging_obj): """Test handler with None websocket messages""" From da73896ec38d7cc5606d1563d87ddd7db31a4d2f Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:43:50 -0400 Subject: [PATCH 012/297] refactor(vertex-live): type the usage helpers without Any The two helpers this branch adds took Sequence[Mapping[str, Any]], which the repo forbids, and only typechecked because Any is compatible with everything. Both now take Mapping[str, object] and the raw *TokensDetails value is narrowed to its mapping entries at each of the three call sites. TypedDicts are the wrong tool here: _merged_modality_totals reads count_key and details_key as runtime strings, and the aggregation deliberately passes unknown keys straight through, so both need a mapping whose keys are not literals. The narrowing is not cosmetic. The handler's only failure path returns no result at all, so a *TokensDetails value that was not a list of objects used to raise while being read and cost the whole session its bill. --- ...tex_ai_live_passthrough_logging_handler.py | 18 +++++++---- .../test_vertex_ai_live_passthrough.py | 30 +++++++++++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index e224f707b02..72e8c3f7657 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -9,7 +9,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( @@ -37,6 +37,11 @@ _AGGREGATED_FIELDS: Final = frozenset( ) +def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]: + """Narrow one turn's ``*TokensDetails`` value to the entries that are actually shaped like one.""" + return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else () + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -68,7 +73,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _resolve_detail_counts( - details: Sequence[Mapping[str, Any]], + details: Sequence[Mapping[str, object]], declared_total: object, ) -> tuple[tuple[str, int], ...]: """ @@ -100,7 +105,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _merged_modality_totals( - snapshots: Sequence[Mapping[str, Any]], + snapshots: Sequence[Mapping[str, object]], count_key: str, details_key: str, ) -> Mapping[str, int]: @@ -109,7 +114,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): tuple( chain.from_iterable( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - snapshot.get(details_key) or [], snapshot.get(count_key) + _detail_entries(snapshot.get(details_key)), snapshot.get(count_key) ) for snapshot in snapshots ) @@ -179,12 +184,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount") + _detail_entries(usage_metadata.get("promptTokensDetails")), usage_metadata.get("promptTokenCount") ) ) candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount") + _detail_entries(usage_metadata.get("candidatesTokensDetails")), + usage_metadata.get("candidatesTokenCount"), ) ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 3b6a548b219..ba9167e2b13 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -381,6 +381,36 @@ class TestVertexAILivePassthroughLoggingHandler: grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" + def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj): + """A ``*TokensDetails`` value that is not a list of objects must not take the session down. + + The handler's only error path returns no result at all, so one odd frame used to throw + while reading it and the whole session billed nothing. The good turns still bill. + """ + turns = self.AUDIO_SESSION[:3] + messages = self._live_messages(turns) + mangled = [dict(message) for message in messages] + mangled[1]["usageMetadata"] = {**mangled[1]["usageMetadata"], "promptTokensDetails": "TEXT"} + + usage = self._session_usage(handler, mock_logging_obj, mangled, self.NATIVE_AUDIO_MODEL) + + surviving = turns[1:] + assert usage.prompt_tokens_details.audio_tokens == sum(turn["prompt"][1] for turn in surviving) + assert usage.prompt_tokens_details.text_tokens == sum(turn["prompt"][0] for turn in surviving) + assert usage.prompt_tokens == sum(sum(turn["prompt"]) for turn in turns), "the totals still cover every turn" + + direct = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 40, + "candidatesTokenCount": 12, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 40}, "AUDIO"], + "candidatesTokensDetails": {"modality": "TEXT", "tokenCount": 12}, + }, + model=self.NATIVE_AUDIO_MODEL, + ) + assert direct.prompt_tokens_details.audio_tokens == 40, "the well-formed entry beside a bad one still counts" + assert direct.completion_tokens == 12 + @pytest.mark.parametrize( "label,prompt_details,candidate_details", [ From b051aa713a3d4c044fb9453935879f09d91eea69 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:54:13 -0400 Subject: [PATCH 013/297] fix(vertex-live): sum tool-use prompt tokens across a session's turns toolUsePromptTokenCount was the one prompt-side total not named in _AGGREGATED_FIELDS, so it rode the unknown-key pass-through and took the first frame's value while promptTokenCount, candidatesTokenCount and totalTokenCount beside it were summed. Live's frames grow over a session, so the first frame is the smallest number in the series and a grounded session under-reported its tool-use tokens by everything after turn one. It is now summed like its three neighbours. This is reporting only, and pricing these tokens is deliberately left out. Google charges tool-use prompt tokens at the input token rate, but generic_cost_per_token reads the input bill out of prompt_tokens_details and only falls back to prompt_tokens when the details carry no text or a cache hit overlaps them. Measured on the native-audio entry with 500 tool-use tokens: adding them to prompt_tokens moves an ordinary Live turn's bill by $0.0000000000, and on a turn with a cache hit it moves it by $0.0002650000 where the tokens are worth $0.0002500000, because it perturbs the cache-overlap correction. Pricing them belongs beside the modality terms in the shared input-cost path, in its own change that fixes the same latent no-op on the ordinary Gemini path. Not verified against a live capture: no Vertex Live session we have captured reported toolUsePromptTokenCount at all, so the summing convention is inferred from the three prompt-side totals that accumulate the same way. --- ...tex_ai_live_passthrough_logging_handler.py | 2 + .../test_vertex_ai_live_passthrough.py | 46 ++++++++++++++----- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 72e8c3f7657..2cd336db3d6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -31,6 +31,7 @@ _AGGREGATED_FIELDS: Final = frozenset( "promptTokenCount", "candidatesTokenCount", "totalTokenCount", + "toolUsePromptTokenCount", "promptTokensDetails", "candidatesTokensDetails", } @@ -159,6 +160,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "toolUsePromptTokenCount": sum(snapshot.get("toolUsePromptTokenCount", 0) for snapshot in snapshots), "promptTokensDetails": [ {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 ], diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index ba9167e2b13..bfff52de673 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -363,22 +363,46 @@ class TestVertexAILivePassthroughLoggingHandler: cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) - def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): - """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + TOOL_USE_PER_TURN = (100, 250, 400) - It is reported, not billed. Nothing in the shared Gemini input-cost path prices - tool-use tokens, and folding them into prompt_tokens here would suppress that - path's cache-overlap correction and raise the bill instead. + def _grounded_messages(self): + """The three-turn session again, with each turn's own toolUsePromptTokenCount attached.""" + messages = self._live_messages(self.AUDIO_SESSION[:3]) + head, turns = messages[0], messages[1:] + return [head] + [ + {**message, "usageMetadata": {**message["usageMetadata"], "toolUsePromptTokenCount": tool_use}} + for message, tool_use in zip(turns, self.TOOL_USE_PER_TURN) + ] + + def test_server_side_tool_use_prompt_tokens_are_summed_over_the_session(self, handler, mock_logging_obj): + """toolUsePromptTokenCount rode the unknown-key pass-through, so it took the first turn only. + + Every other total beside it is summed across the session, and the first turn is the + smallest number in the series, so a grounded session logged far fewer tool-use tokens + than it used. This session's turns are deliberately distinct, so 750 can only come from + summing: first-turn selection gives 100, last-turn or max gives 400. """ - messages = self._live_messages(self.AUDIO_SESSION[:1]) - grounded = [dict(message) for message in messages] - grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + grounded = self._grounded_messages() usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) - assert usage.prompt_tokens_details.tool_use_tokens == 500 + assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) - plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) - grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): + """Deliberate boundary: these tokens are reported here, and priced nowhere. + + generic_cost_per_token reads the input bill out of prompt_tokens_details, and falls + back to prompt_tokens only when the details carry no text or a cache hit overlaps them, + so adding tool-use tokens to prompt_tokens is worth nothing on an ordinary Live turn and + over-charges against the cache-overlap correction when it is not. Pricing them belongs + in the shared input-cost path, beside the modality terms that already read the details. + """ + turns = self.AUDIO_SESSION[:3] + plain_cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost( + handler, mock_logging_obj, self._grounded_messages(), self.NATIVE_AUDIO_MODEL + ) + + assert plain_cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj): From 22e7c7a5338a96009607c612804fe90c0dad9a1a Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 06:25:07 -0400 Subject: [PATCH 014/297] test(vertex-live): type the session helpers in the Live passthrough tests The four session helpers this PR added were unannotated. Typing them needs a name for the (text, audio) pair each turn carries, so _LiveTurn is a TypedDict rather than a Mapping union that would leave sum() over a prompt pair ill-typed, and AUDIO_SESSION is declared with it. The message list reuses list[dict[str, object]], the annotation the passthrough already uses where it collects those messages --- .../test_vertex_ai_live_passthrough.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index bfff52de673..b4d0a6c06e5 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -6,12 +6,14 @@ including the logging handler, cost tracking, and WebSocket message processing. """ import json +from collections.abc import Sequence from datetime import datetime from unittest.mock import AsyncMock, Mock, patch, MagicMock from typing import Dict, List, Any, Optional import pytest import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict # Add the parent directory to the system path @@ -22,10 +24,16 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.utils import LlmProviders +from litellm.types.utils import LlmProviders, Usage from litellm.proxy._types import UserAPIKeyAuth +class _LiveTurn(TypedDict): + prompt: ReadOnly[tuple[int, int]] + candidates: ReadOnly[tuple[int, int]] + candidate_audio_token_count_missing: NotRequired[ReadOnly[bool]] + + class TestVertexAILivePassthroughLoggingHandler: """Test the Vertex AI Live Passthrough Logging Handler""" @@ -257,7 +265,7 @@ class TestVertexAILivePassthroughLoggingHandler: # window, so the prompt side repeats the accumulated audio while the candidates side reports # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is # the shape Live really emits at the end of a spoken answer. - AUDIO_SESSION = ( + AUDIO_SESSION: tuple[_LiveTurn, ...] = ( {"prompt": (14, 122), "candidates": (8, 20)}, {"prompt": (21, 182), "candidates": (5, 50)}, {"prompt": (24, 203), "candidates": (13, 27)}, @@ -265,7 +273,7 @@ class TestVertexAILivePassthroughLoggingHandler: ) @staticmethod - def _live_messages(turns): + def _live_messages(turns: Sequence[_LiveTurn]) -> list[dict[str, object]]: """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits.""" return [{"type": "session.created", "session": {"id": "s"}}] + [ { @@ -292,7 +300,12 @@ class TestVertexAILivePassthroughLoggingHandler: ] @staticmethod - def _session_usage(handler, mock_logging_obj, messages, model): + def _session_usage( + handler: VertexAILivePassthroughLoggingHandler, + mock_logging_obj: MagicMock, + messages: list[dict[str, object]], + model: str, + ) -> Usage: result = handler.vertex_ai_live_passthrough_handler( websocket_messages=messages, logging_obj=mock_logging_obj, @@ -306,7 +319,13 @@ class TestVertexAILivePassthroughLoggingHandler: return result["result"].usage @classmethod - def _session_cost(cls, handler, mock_logging_obj, messages, model): + def _session_cost( + cls, + handler: VertexAILivePassthroughLoggingHandler, + mock_logging_obj: MagicMock, + messages: list[dict[str, object]], + model: str, + ) -> float: from litellm.cost_calculator import completion_cost from litellm.types.utils import ModelResponse @@ -321,7 +340,7 @@ class TestVertexAILivePassthroughLoggingHandler: ) @classmethod - def _expected_session_cost(cls, turns): + def _expected_session_cost(cls, turns: Sequence[_LiveTurn]) -> float: from litellm.utils import get_model_info info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai") From b1262b05f470f550ea0a2d879449f4fd966c4a1d Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 22:05:02 -0400 Subject: [PATCH 015/297] fix(gemini-live): count grounding requests so Live sessions carry their query fee Live reports grounding in the server frames and never in usageMetadata, so nothing set the counter the cost path reads and the per-query charge was missing from every grounded session. Google bills a grounded Live prompt on top of its tokens, and that fee dwarfs the token cost, so a non-zero spend check could never catch it. Both Live surfaces now read serverContent.groundingMetadata where they build usage, and reuse the chat path's own classifier so web search and Maps keep their separate SKUs rather than being counted together. Separately, a client sending turn_detection: null reached a membership test against None and took the session down with no traceback, while the branch immediately above already guards for it. Live emits grounding and usage on the same frame, verified against Vertex directly, so the realtime counter is set where usage is built. (cherry picked from commit c997436be34beb2e84f8468286b8016c195eca92) (cherry picked from commit 26c8d4822fc1c5c44fe8f72f4f57473a2ce1acbf) --- .../llms/gemini/realtime/transformation.py | 20 +++- ...tex_ai_live_passthrough_logging_handler.py | 32 ++++++- .../test_vertex_ai_live_passthrough.py | 68 +++++++++++++ .../test_gemini_realtime_transformation.py | 95 ++++++++++++++++++- 4 files changed, 212 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index c92af7de145..79985569c5f 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -115,6 +115,19 @@ def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup return envelope.get("setup", empty_setup) +def _grounding_metadata_from_frame(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + """Read ``serverContent.groundingMetadata`` off the frame that carries the turn's usage. + + Live reports grounding in the server frames rather than in ``usageMetadata``, and it emits both + on the same frame, so the per-query charge is countable at the point usage is built. + """ + server_content: Final = frame.get("serverContent") + if not isinstance(server_content, Mapping): + return () + metadata: Final = server_content.get("groundingMetadata") + return (metadata,) if isinstance(metadata, Mapping) else () + + # Google bills Live transcription at an estimated 25 audio tokens/sec of input and # 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing). GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25 @@ -323,7 +336,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} - elif key == "turn_detection": + elif key == "turn_detection" and value is not None: value_typed = cast(OpenAIRealtimeTurnDetection, value) if ( isinstance(value_typed, dict) @@ -1049,6 +1062,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): {**cast(dict, message), "usageMetadata": resolved_usage_metadata}, ), ) + grounding_metadata: Final = _grounding_metadata_from_frame(message) + if grounding_metadata: + VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet + _chat_completion_usage, grounding_metadata + ) else: _chat_completion_usage = get_empty_usage() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 2cd336db3d6..0e2eb60704d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -43,6 +43,23 @@ def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]: return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else () +def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[str, object], ...]: + """Collect every ``serverContent.groundingMetadata`` a session emitted. + + Live reports grounding in the server frames, never in ``usageMetadata``, so the per-query + charge has to be counted here rather than derived from the token totals. + """ + return tuple( + metadata + for message in websocket_messages + if isinstance(message, Mapping) + for server_content in (message.get("serverContent"),) + if isinstance(server_content, Mapping) + for metadata in (server_content.get("groundingMetadata"),) + if isinstance(metadata, Mapping) + ) + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -173,6 +190,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): def _create_usage_object_from_metadata( usage_metadata: dict, model: str, + grounding_metadata: Sequence[Mapping[str, object]] = (), ) -> Usage: """ Create a LiteLLM Usage object from Live API usage metadata. @@ -180,6 +198,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Args: usage_metadata: Usage metadata from the Live API response model: The model name + grounding_metadata: Every ``serverContent.groundingMetadata`` the session emitted, so + Search and Maps grounding carry their per-query charge Returns: LiteLLM Usage object @@ -199,7 +219,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) - return Usage( + usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), @@ -217,6 +237,15 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): video_tokens=candidates_by_modality.get("VIDEO"), ), ) + if grounding_metadata: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet + usage, grounding_metadata + ) + return usage def vertex_ai_live_passthrough_handler( self, @@ -264,6 +293,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Create Usage object for standard LiteLLM logging usage: Final = self._create_usage_object_from_metadata( usage_metadata=usage_metadata, + grounding_metadata=_grounding_metadata(websocket_messages), model=model, ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index b4d0a6c06e5..1815ff134aa 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -406,6 +406,74 @@ class TestVertexAILivePassthroughLoggingHandler: usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) + @staticmethod + def _grounding_frame(metadata: dict[str, object]) -> dict[str, object]: + """One server frame carrying grounding metadata, the way Live reports it.""" + return {"type": "response.done", "serverContent": {"groundingMetadata": metadata}} + + def test_web_grounding_is_counted_so_it_can_be_billed(self, handler, mock_logging_obj): + """Live reports grounding in the server frames and never in usageMetadata. + + Nothing read those frames, so web_search_requests stayed unset and the cost path's only + trigger for the per-query grounding charge never fired. Google bills a grounded Live + prompt on top of its tokens, so the whole fee was missing from the bill. + """ + messages = [ + self._grounding_frame( + { + "webSearchQueries": ["who won the 2026 world cup final"], + "groundingChunks": [{"web": {"uri": "https://example.com"}}], + } + ), + *self._live_messages(self.AUDIO_SESSION[:1]), + ] + + usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + + assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + + def test_maps_grounding_is_counted_under_its_own_sku(self, handler, mock_logging_obj): + """Maps grounding is a separate SKU from web search, so it needs its own counter. + + A maps-only turn carries grounding chunks but no webSearchQueries, so counting queries + alone would report nothing and bill nothing. + """ + messages = [ + self._grounding_frame({"groundingChunks": [{"maps": {"placeId": "abc123"}}]}), + *self._live_messages(self.AUDIO_SESSION[:1]), + ] + + usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + + assert usage.prompt_tokens_details.google_maps_grounding_requests == 1 + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + + def test_an_ungrounded_session_reports_no_grounding(self, handler, mock_logging_obj): + """The counters must stay absent when no tool ran, or every session pays a grounding fee.""" + usage = self._session_usage( + handler, mock_logging_obj, self._live_messages(self.AUDIO_SESSION[:1]), self.NATIVE_AUDIO_MODEL + ) + + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + + def test_grounding_adds_its_query_fee_to_the_session_bill(self, handler, mock_logging_obj): + """The counter only matters if it reaches the bill, so assert against the cost, not the field. + + Same tokens either way: the difference between the two sessions is the grounding fee alone. + """ + turns = self.AUDIO_SESSION[:1] + plain = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded = self._session_cost( + handler, + mock_logging_obj, + [self._grounding_frame({"webSearchQueries": ["q"]}), *self._live_messages(turns)], + self.NATIVE_AUDIO_MODEL, + ) + + assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded" + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 2b3b6343fad..ee984cc7e1f 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,11 +1,16 @@ import json -from unittest.mock import MagicMock +from collections.abc import Mapping +from typing import cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig +from litellm.types.llms.gemini import BidiGenerateContentServerMessage +from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents +from litellm.types.utils import Usage def test_gemini_realtime_transformation_session_created(): @@ -2178,3 +2183,91 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra } assert usage == expected assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None + + +def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]: + """One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them.""" + from typing import Final + + server_content: Final = { + "turnComplete": True, + **({} if grounding_metadata is None else {"groundingMetadata": grounding_metadata}), + } + return { + "serverContent": server_content, + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 157, + "totalTokenCount": 176, + "promptTokensDetails": ({"modality": "TEXT", "tokenCount": 19},), + "candidatesTokensDetails": ({"modality": "AUDIO", "tokenCount": 157},), + }, + } + + +def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage: + """Capture the chat-completion Usage transform_response_done_event builds, before it is bridged. + + The Usage object is local to the method, so the bridge call is the only place it is observable. + """ + from typing import Final + + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes + original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage + + def _spy(usage: Usage) -> object: + captured.append(usage) + return original(usage) + + config: Final = GeminiRealtimeConfig() + with patch.object( + LiteLLMCompletionResponsesConfig, + "_transform_chat_completion_usage_to_responses_usage", + staticmethod(_spy), + ): + config.transform_response_done_event( + message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict + BidiGenerateContentServerMessage, message + ), + current_response_id="resp_grounding", + current_conversation_id="conv_grounding", + output_items=None, + ) + assert captured, "response.done must build a Usage object" + return captured[0] + + +def test_gemini_realtime_response_done_counts_web_grounding(): + """Regression: Live reports grounding in the server frames and never in usageMetadata. + + Nothing read those frames on the realtime path, so web_search_requests stayed unset and the + cost path's only trigger for Google's per-query grounding charge never fired. + + Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is + built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of + detail fields and drops the rest, so the counter does not reach response.done yet. Widening + that bridge is a separate change; do not read this test as proving end-to-end billing. + """ + usage = _usage_built_for_response_done( + _grounded_live_frame( + { + "webSearchQueries": ["who won the 2026 world cup final"], + "groundingChunks": [{"web": {"uri": "https://example.com"}}], + } + ) + ) + + assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" + assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it" + + +def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran(): + """The counter must stay unset on an ordinary turn, or every session pays a grounding fee.""" + usage = _usage_built_for_response_done(_grounded_live_frame(None)) + + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None From 9580b89bb16d9946152aa23a7a4d7d139700d403 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 23:26:31 -0400 Subject: [PATCH 016/297] fix(vertex-live): resolve a Live setup model before logging reads it A client that named a bare gateway alias logged the session as "unknown" and billed nothing, because the model was read off the raw setup frame and the extractor only yields a name when the string already contains "/models/". The rewriter qualifies that same model a few lines later for the upstream, so the supported client form, an alias, was the one that went unbilled. Resolving through the rewriter first means the real model reaches the logging object, and from there the cost map. A route with no rewriter, which is every non-Live passthrough, hands the frame over untouched. (cherry picked from commit 573982803df612fd94144e2e06dd647f8530d4e8) --- .../pass_through_endpoints.py | 23 ++++- .../test_pass_through_endpoints.py | 94 +++++++++++++++++++ 2 files changed, 116 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index ea4ede7e513..5a2f9c391a2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload +def _resolved_vertex_live_setup( + setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None +) -> Mapping[str, object]: + """ + Give the model extractor the same fully qualified path the upstream will receive. + + Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before + it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw + frame logs the session as ``unknown`` at no cost, which is precisely the supported client form + """ + setup_model: Final = setup_data.get("model") + if setup_model_rewriter is None or not isinstance(setup_model, str): + return setup_data + return {**setup_data, "model": setup_model_rewriter(setup_model)} + + def _truncated_close_reason(reason: str) -> str: """ Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character @@ -2314,7 +2330,12 @@ async def websocket_passthrough_request( setup_data, ) if isinstance(setup_data, dict) and "model" in setup_data: - extracted_model = _extract_model_from_vertex_ai_setup(setup_data) + # Resolve the alias first: a client may name a bare gateway model, + # which carries no "/models/" for the extractor to read, so reading + # the raw frame leaves the session logged as "unknown" and unbilled. + extracted_model = _extract_model_from_vertex_ai_setup( + _resolved_vertex_live_setup(setup_data, setup_model_rewriter) + ) if extracted_model: kwargs["model"] = extracted_model kwargs["custom_llm_provider"] = "vertex_ai-language-models" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d57bed430c1..11066d4ed38 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5030,6 +5030,100 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model(): assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" +@pytest.mark.parametrize( + "setup_model", + ["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"], +) +def test_vertex_live_setup_model_resolves_before_extraction(setup_model): + """A bare gateway alias left the session logged as ``unknown`` at zero cost. + + The model was read off the raw client frame, and the extractor only yields a name when the string + already contains ``/models/``. The rewriter qualifies it a few lines later for the upstream, so a + client that addressed the gateway the documented way, by alias, logged no model and therefore + resolved no cost-map entry. Resolving first is what puts the real name on the logging object. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _extract_model_from_vertex_ai_setup, + _resolved_vertex_live_setup, + ) + + rewriter = _build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", vertex_location="global", llm_router=None + ) + setup_data = {"model": setup_model} + + resolved = _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, rewriter)) + + assert resolved == "gemini-live-2.5-flash", "an unresolved setup model logs the session as 'unknown'" + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_a_bare_alias_setup_model(): + """End to end through the relay: a bare alias must reach the logging object as a real model name. + + This is the call-site half of the fix. The helper tests above pass even if extraction moves back + before the rewrite, so this one drives the real websocket relay and asserts on what got logged, + which is the name the cost map is looked up by. An unbilled session logs ``unknown``. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + + upstream_ws = RecordingUpstreamWebSocket() + setup_frame = json.dumps({"setup": {"model": "gemini-live-2.5-flash"}}) + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "text": setup_frame}, + {"type": "websocket.disconnect"}, + ] + ) + ) + built = [] + real_logging = litellm.litellm_core_utils.litellm_logging.Logging + + def _capture(*args, **kwargs): + obj = real_logging(*args, **kwargs) + built.append(obj) + return obj + + with _patched_websocket_passthrough_environment(upstream_ws): + with patch("litellm.litellm_core_utils.litellm_logging.Logging", side_effect=_capture): + await websocket_passthrough_request( + websocket=websocket, + target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + setup_model_rewriter=_build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", vertex_location="global", llm_router=None + ), + ) + + assert built, "the relay should have built a logging object" + assert built[0].model == "gemini-live-2.5-flash", "a bare alias must not log as 'unknown'" + + +def test_vertex_live_setup_resolution_is_inert_without_a_rewriter(): + """Non-Live passthrough routes pass no rewriter, so the frame must be handed over untouched.""" + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _extract_model_from_vertex_ai_setup, + _resolved_vertex_live_setup, + ) + + setup_data = {"model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"} + + assert _resolved_vertex_live_setup(setup_data, None) is setup_data + assert _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, None)) == ( + "gemini-live-2.5-flash" + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("rcvd_close", [None, "abnormal", "no_status"]) async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rcvd_close): From 6d71e3385b8c2ed0db61cfc1991de5614de15d83 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 22:12:46 -0400 Subject: [PATCH 017/297] fix(cost): carry grounding counters through the Responses usage bridge A realtime session's usage is rebuilt from its own response.done event, so a counter that does not survive the round trip is invisible to the cost path. Both directions copied a fixed allow-list, which meant a grounded Gemini Live session reported its query on the Usage object and then lost it before anything could bill it. Gemini reads the grounding counters off the input token details while Anthropic reads its own server_tool_use field, so carrying these two cannot move an Anthropic bill. Absent counters stay absent, so no provider starts paying a fee it did not incur. (cherry picked from commit ffd6c723e2a65dfff1860a913c34e41bc72ff11f) (cherry picked from commit 1590822f6893c57086b72965963d23f4f367b424) --- .../litellm_completion_transformation/transformation.py | 8 ++++++++ litellm/responses/utils.py | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..2b011abf0e2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2667,6 +2667,14 @@ class LiteLLMCompletionResponsesConfig: if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None: input_details_dict["audio_tokens"] = prompt_details.audio_tokens + # The cost path reads the grounding counters off the input details, and a realtime + # session's usage is rebuilt from its own response.done, so dropping them here bills + # no per-query grounding fee at all. + for counter in ("web_search_requests", "google_maps_grounding_requests"): + counter_value = getattr(prompt_details, counter, None) + if counter_value is not None: + input_details_dict[counter] = counter_value + cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) or getattr( prompt_details, "cache_creation_tokens", None ) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 540d492beec..7e0acd14b9c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1137,6 +1137,12 @@ class ResponseAPILoggingUtils: text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), + web_search_requests=getattr( + response_api_usage.input_tokens_details, "web_search_requests", None + ), + google_maps_grounding_requests=getattr( + response_api_usage.input_tokens_details, "google_maps_grounding_requests", None + ), ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None output_tokens_details: Final[OutputTokensDetails | None] = getattr( From 6c83484065ac1b6fbed79bb9d0acea97f50a0915 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Tue, 8 Sep 2026 00:59:28 -0400 Subject: [PATCH 018/297] chore(vertex-live): drop a comment that restated the helper's docstring The call site repeated _resolved_vertex_live_setup's own docstring almost verbatim, which is the duplication the repo's comment rule exists to prevent. --- litellm/proxy/pass_through_endpoints/pass_through_endpoints.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a2f9c391a2..b66c295d1aa 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2330,9 +2330,6 @@ async def websocket_passthrough_request( setup_data, ) if isinstance(setup_data, dict) and "model" in setup_data: - # Resolve the alias first: a client may name a bare gateway model, - # which carries no "/models/" for the extractor to read, so reading - # the raw frame leaves the session logged as "unknown" and unbilled. extracted_model = _extract_model_from_vertex_ai_setup( _resolved_vertex_live_setup(setup_data, setup_model_rewriter) ) From 4e51ff8a6dbcc6a544891f8e64b887c8cd6a8ae2 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 9 Sep 2026 00:00:46 -0400 Subject: [PATCH 019/297] style(cost): collapse a getattr call that fits the line limit --- litellm/responses/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 7e0acd14b9c..b2844834860 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1137,9 +1137,7 @@ class ResponseAPILoggingUtils: text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), - web_search_requests=getattr( - response_api_usage.input_tokens_details, "web_search_requests", None - ), + web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( response_api_usage.input_tokens_details, "google_maps_grounding_requests", None ), From 228d87db6341a87bd9ae2bd1ebd76ae5e08e1a08 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 9 Sep 2026 00:19:12 -0400 Subject: [PATCH 020/297] test(gemini-live): assert grounding counters on the emitted response.done event Replaces a patch.object spy on a static method with a read of the public return value, which also covers the usage bridge the spy ran ahead of. --- .../test_gemini_realtime_transformation.py | 66 +++++++------------ 1 file changed, 22 insertions(+), 44 deletions(-) diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index ee984cc7e1f..3eb4a70ee15 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,7 +1,7 @@ import json from collections.abc import Mapping from typing import cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -9,8 +9,6 @@ import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig from litellm.types.llms.gemini import BidiGenerateContentServerMessage -from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents -from litellm.types.utils import Usage def test_gemini_realtime_transformation_session_created(): @@ -2205,40 +2203,22 @@ def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Map } -def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage: - """Capture the chat-completion Usage transform_response_done_event builds, before it is bridged. - - The Usage object is local to the method, so the bridge call is the only place it is observable. - """ +def _response_done_input_details(message: Mapping[str, object]) -> Mapping[str, object]: + """The ``input_tokens_details`` a ``response.done`` event carries, read off the emitted event.""" from typing import Final - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes - original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage - - def _spy(usage: Usage) -> object: - captured.append(usage) - return original(usage) - config: Final = GeminiRealtimeConfig() - with patch.object( - LiteLLMCompletionResponsesConfig, - "_transform_chat_completion_usage_to_responses_usage", - staticmethod(_spy), - ): - config.transform_response_done_event( - message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict - BidiGenerateContentServerMessage, message - ), - current_response_id="resp_grounding", - current_conversation_id="conv_grounding", - output_items=None, - ) - assert captured, "response.done must build a Usage object" - return captured[0] + event: Final = config.transform_response_done_event( + message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict + BidiGenerateContentServerMessage, message + ), + current_response_id="resp_grounding", + current_conversation_id="conv_grounding", + output_items=None, + ) + usage: Final = event["response"]["usage"] + assert usage, "response.done must carry a usage object" + return usage.get("input_tokens_details") or {} def test_gemini_realtime_response_done_counts_web_grounding(): @@ -2247,12 +2227,10 @@ def test_gemini_realtime_response_done_counts_web_grounding(): Nothing read those frames on the realtime path, so web_search_requests stayed unset and the cost path's only trigger for Google's per-query grounding charge never fired. - Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is - built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of - detail fields and drops the rest, so the counter does not reach response.done yet. Widening - that bridge is a separate change; do not read this test as proving end-to-end billing. + The counter is read off the emitted event, which is what the cost path is handed, so this covers + the grounding read and the usage bridge that carries it together """ - usage = _usage_built_for_response_done( + input_details = _response_done_input_details( _grounded_live_frame( { "webSearchQueries": ["who won the 2026 world cup final"], @@ -2261,13 +2239,13 @@ def test_gemini_realtime_response_done_counts_web_grounding(): ) ) - assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" - assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it" + assert input_details.get("web_search_requests") == 1, "a grounded turn must report its query" + assert input_details.get("text_tokens") == 19, "the modality breakdown must survive alongside it" def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran(): """The counter must stay unset on an ordinary turn, or every session pays a grounding fee.""" - usage = _usage_built_for_response_done(_grounded_live_frame(None)) + input_details = _response_done_input_details(_grounded_live_frame(None)) - assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None - assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + assert input_details.get("web_search_requests") is None + assert input_details.get("google_maps_grounding_requests") is None From 36c1e5e17d1326f1a8f3dc7b25e86a69349d53e7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:34:31 -0700 Subject: [PATCH 021/297] refactor(compression): build the protected index set without mutation --- litellm/compression/compress.py | 33 ++++--------------- .../guardrail_hooks/test_headroom.py | 26 --------------- 2 files changed, 7 insertions(+), 52 deletions(-) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 62b05a4938f..c79e6aed57a 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -206,13 +206,6 @@ def _extract_anthropic_tool_exchange_spans( def _message_has_cache_control(message: Mapping[str, object]) -> bool: - """True if ``message`` carries an Anthropic ``cache_control`` breakpoint. - - A breakpoint can sit directly on the message dict, or on any part of a - list-of-parts ``content`` (the shape Anthropic's own messages use). Either - placement pins the provider's KV-cache prefix to this row's exact bytes, so - either placement must protect the row the same way. - """ if message.get("cache_control") is not None: return True content: Final = message.get("content") @@ -231,28 +224,16 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int The last user message is what the model is being asked to act on right now, so compressing it replaces the live instruction with a marker. Compression - guardrails share this policy; see the Headroom guardrail. - - A cache_control breakpoint pins the provider's prompt-cache prefix to that - row's exact bytes. Rewriting the row (even leaving the marker in place) - changes those bytes, so the next request misses the cache it thinks it is - reusing and silently pays a cache write instead of a cache read. This is - not limited to the last user/assistant row: a marker several turns back - (e.g. on a large cached tool result) needs the same protection. + guardrails share this policy; see the Headroom guardrail. A cache_control + breakpoint pins the provider's prompt-cache prefix to that row's exact + bytes, so rewriting a marked row anywhere in history turns the next + request's cache read into a cache write. """ system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system") last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:] - last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] - cache_control_indices: Final = tuple( - index for index, msg in enumerate(messages) if _message_has_cache_control(msg) - ) - seen: Final[set[int]] = set() - ordered: Final[list[int]] = [] - for index in system_indices + last_user + last_assistant + cache_control_indices: - if index not in seen: - seen.add(index) - ordered.append(index) - return tuple(ordered) + assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant") + cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg)) + return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices)) def _combine_scores( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 5cd42bd3f83..d4531398ba1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1795,14 +1795,6 @@ PARTS_MESSAGES = [ ], }, { - # No cache_control here on purpose: this row exercises the general - # multi-part flatten/merge mechanics (shared with compresr). A row - # carrying its own cache_control is a different, dedicated case -- - # see test_mid_history_cache_control_row_is_never_sent_for_compression - # (#39519): get_protected_indices withholds it from /v1/compress - # entirely rather than letting it be rewritten and re-merged, because - # rewriting the bytes under a live breakpoint busts the cache the - # marker is supposed to preserve. "role": "user", "content": [ {"type": "text", "text": "Earlier turn."}, @@ -1895,14 +1887,6 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( messages = result["structured_messages"] history_content = messages[1]["content"] - # Rewritten all-text row collapses to one part carrying the rewritten text. - # This fixture row carries no cache_control (see PARTS_MESSAGES): the - # last-declared-breakpoint-survives-the-merge behavior is a property of - # merge_rewritten_text_parts and is covered directly by compresr's - # test_all_text_row_merges_and_keeps_last_cache_control, since a - # cache_control-marked row never reaches this merge path through Headroom - # at all -- get_protected_indices withholds it before compression runs - # (see test_mid_history_cache_control_row_is_never_sent_for_compression). assert isinstance(history_content, list) assert len(history_content) == 1 assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac" @@ -2530,15 +2514,6 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert messages[3] == compressed_history[1] -# --------------------------------------------------------------------------- -# #39519: a mid-history row carrying its own Anthropic cache_control marker -# (e.g. a large tool result the client already cached several turns back) was -# still sent to /v1/compress and rewritten. It came back byte-different but -# kept its marker, so the next request's cache read silently became a cache -# write. get_protected_indices() now protects any cache_control-marked row, -# not just system/last-user/last-assistant, so it must never reach the wire. -# --------------------------------------------------------------------------- - CACHE_MARKED_HISTORY_MESSAGES = [ {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, {"role": "user", "content": "old question " + "Q" * 5000}, @@ -2565,7 +2540,6 @@ async def test_mid_history_cache_control_row_is_never_sent_for_compression(guard cached_row = CACHE_MARKED_HISTORY_MESSAGES[3] assert cached_row not in wire assert not any(row.get("tool_call_id") == "old_1" for row in wire) - # Byte-identical, marker intact -- the next request's cache read survives. assert result["structured_messages"][3] == cached_row From 492251c7bc7bd76baf25512e4c43421efbcff799 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 10 Sep 2026 07:32:54 +0000 Subject: [PATCH 022/297] fix(otel): cap per-index OpenInference message attributes span-wide OpenInferenceMapper spelled every captured prompt and response message out as two indexed attributes with no bound. A few dozen turns overran the OTel SDK's 128-attribute span limit, which evicts oldest first, so the gen_ai.* model, provider, usage, cost and finish reason written before it were what got dropped. Both directions now share one MAX_MESSAGE_ATTRS_PER_SPAN ceiling, the response keeps at least half of it, and input.value / output.value still carry the complete conversation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../otel/mappers/openinference.py | 26 ++- litellm/integrations/otel/mappers/utils.py | 12 ++ .../integrations/otel/test_otel_v2_emitter.py | 148 ++++++++++++++++++ 3 files changed, 181 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index 0ba45170b8e..1e2dbf6974d 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -12,6 +12,7 @@ from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData from litellm.integrations.otel.mappers.utils import ( + MAX_MESSAGE_ATTRS_PER_SPAN, MAX_TOOL_DEFINITION_ATTRS_PER_SPAN, collect, drop_none, @@ -26,6 +27,8 @@ from litellm.integrations.otel.model.payloads import ( ToolDefinition, ) +_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2 + class OpenInferenceMapper: """Emits OpenInference attributes for LLM_CALL spans. @@ -84,22 +87,35 @@ class OpenInferenceMapper: return {} def _llm_call(self, data: LLMCallSpanData) -> AttributeMap: + outputs: Final = output_messages(data) + indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs)) return { **collect(self._LLM_CALL_ATTRS, data), **collect(self._BLOB_ATTRS, data), - **self._messages("llm.input_messages", "input.value", data.messages_in), - **self._messages("llm.output_messages", "output.value", output_messages(data)), + **self._messages("llm.input_messages", "input.value", data.messages_in, indexed_in), + **self._messages("llm.output_messages", "output.value", outputs, indexed_out), **self._tools(data), } @staticmethod - def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap: - """Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob.""" + def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]: + """How many prompt and response messages get per-index attributes. + + Both directions share one span-wide allowance. The response is reserved at + least half of it, so a long prompt can never push the completion off the + span, and the prompt takes whatever the response leaves unused. + """ + indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs)) + return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out + + @staticmethod + def _messages(prefix: str, value_key: str, messages: Sequence[object], indexed: int) -> AttributeMap: + """``{prefix}.{idx}.message.*`` keys for the leading ``indexed`` messages + the ``value_key`` blob of all.""" parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages] attrs: Final = drop_none( { key: value - for idx, (role, content) in enumerate(parsed) + for idx, (role, content) in enumerate(parsed[:indexed]) for key, value in ( ( f"{prefix}.{idx}.message.role", diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index d45dca782b2..8d21b774319 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -32,6 +32,18 @@ core telemetry no matter how many vocabularies are configured. """ +MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 4 +"""Span-wide ceiling on attributes spent spelling out chat messages per index. + +A conversation is the other unbounded family: two attributes per message, for +the prompt and the response alike, on the same span. Past a few dozen turns the +family alone exceeds the span attribute limit and evicts the core telemetry +written before it. The ceiling covers both directions together, since a budget +handed to each direction separately doubles. The complete conversation still +rides the JSON blob attributes; only the per-index convenience keys are capped. +""" + + def tool_attr_budget(vocabularies: int) -> int: """Split the span-wide tool-definition ceiling across active vocabularies.""" return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index b1b1b62c820..a417bd62124 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -1,6 +1,8 @@ """Golden tests for the OTel v2 engine: span shape, kinds, semconv attributes, legacy dual-emit, hierarchy, error status, and idempotency. Needs the OTel SDK.""" +import json + import pytest pytest.importorskip("opentelemetry") @@ -18,6 +20,7 @@ from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402 from litellm.integrations.otel.emitter import stamp_error # noqa: E402 from litellm.integrations.otel.mappers.utils import ( # noqa: E402 + MAX_MESSAGE_ATTRS_PER_SPAN, MAX_TOOL_DEFINITION_ATTRS_PER_SPAN, ) from litellm.integrations.otel.model.payloads import ( # noqa: E402 @@ -440,3 +443,148 @@ def test_vendor_tool_definitions_are_truncated_not_dropped(): assert a["llm.tools.0.tool.name"] == "tool_0" assert a["llm.tools.0.tool.json_schema"] assert "llm.tools.126.tool.name" not in a + + +def _conversation_payload(turns, choices=1, **overrides): + """A ``turns``-message chat with ``choices`` response choices, content-bearing.""" + return _payload( + messages=[{"role": ("user", "assistant")[i % 2], "content": f"turn {i}"} for i in range(turns)], + response={ + "id": "resp_1", + "model": "gpt-4o-2024", + "choices": [ + {"finish_reason": "stop", "message": {"role": "assistant", "content": f"reply {i}"}} + for i in range(choices) + ], + }, + **overrides, + ) + + +def _conversation_span(mapper_names, payload): + """The exported LLM-call span for ``payload`` with content capture on.""" + cfg = OpenTelemetryV2Config( + exporter="in_memory", + mapper_names=list(mapper_names), + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + engine.emit( + SpanRole.LLM_CALL, + LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True), + ) + (span,) = exporter.get_finished_spans() + return span + + +def _indexed_message_count(attributes, prefix): + return len({key.split(".")[2] for key in attributes if key.startswith(f"{prefix}.")}) + + +@pytest.mark.parametrize("turns", [60, 200]) +def test_long_conversation_does_not_evict_core_attributes(turns): + """Per-message OpenInference attributes must never crowd core telemetry off the span. + + With content capture on, the OpenInference vocabulary spells every prompt and + response message out as two per-index attributes. A few dozen turns overruns + the OTel SDK's 128-attribute span limit, which evicts oldest-first, so the + ``gen_ai.*`` set written before it is what disappears. + """ + span = _conversation_span(["genai", "openinference"], _conversation_payload(turns)) + a = span.attributes + + assert span.dropped_attributes == 0 + assert a[GenAI.REQUEST_MODEL] == "gpt-4o" + assert a[GenAI.PROVIDER_NAME] == "openai" + assert a[GenAI.USAGE_INPUT_TOKENS] == 10 + assert a[GenAI.USAGE_OUTPUT_TOKENS] == 5 + assert a[GenAI.RESPONSE_FINISH_REASONS] == ("stop",) + assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 + + assert a["llm.input_messages.0.message.content"] == "turn 0" + assert a["llm.output_messages.0.message.content"] == "reply 0" + assert f"llm.input_messages.{turns - 1}.message.role" not in a + assert len(json.loads(a["input.value"])) == turns + assert len(json.loads(a["output.value"])) == 1 + assert len(json.loads(a[GenAI.INPUT_MESSAGES])) == turns + + +def test_short_conversation_keeps_every_message_indexed(): + """Below the cap nothing is truncated in either direction.""" + a = _conversation_span(["genai", "openinference"], _conversation_payload(4, choices=2)).attributes + for idx in range(4): + assert a[f"llm.input_messages.{idx}.message.content"] == f"turn {idx}" + for idx in range(2): + assert a[f"llm.output_messages.{idx}.message.content"] == f"reply {idx}" + + +def test_message_cap_is_shared_across_input_and_output(): + """One span-wide allowance covers both directions, and the response always keeps a share. + + A long prompt takes what a single reply leaves over, and a many-choice reply + cannot take the whole allowance away from the prompt either. + """ + long_prompt = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=1)).attributes + many_choices = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=20)).attributes + + single_reply_indexed = _indexed_message_count(long_prompt, "llm.output_messages") + assert single_reply_indexed == 1 + assert _indexed_message_count(long_prompt, "llm.input_messages") + single_reply_indexed == ( + MAX_MESSAGE_ATTRS_PER_SPAN // 2 + ) + + assert _indexed_message_count(many_choices, "llm.input_messages") > 0 + assert _indexed_message_count(many_choices, "llm.output_messages") > single_reply_indexed + assert _indexed_message_count(many_choices, "llm.input_messages") + _indexed_message_count( + many_choices, "llm.output_messages" + ) == (MAX_MESSAGE_ATTRS_PER_SPAN // 2) + + +def test_fully_populated_arize_span_stays_within_the_attribute_limit(): + """Every capped family maxed at once still leaves the whole core intact. + + The Arize / Phoenix composition (``genai`` + ``openinference`` + ``legacy``) + with every request parameter, every cost component, a hundred-plus tools, a + two-hundred-turn prompt and twenty choices is the worst case the two + span-wide ceilings have to absorb together. + """ + payload = _conversation_payload( + 200, + choices=20, + stream=True, + model_parameters={ + **_tools_payload(127)["model_parameters"], + "top_p": 0.9, + "frequency_penalty": 0.1, + "presence_penalty": 0.1, + "seed": 7, + "stop": ["\n"], + }, + cost_breakdown={ + key: 0.001 + for key in ( + "input_cost", + "output_cost", + "cache_read_cost", + "cache_creation_cost", + "tool_usage_cost", + "original_cost", + "discount_amount", + "discount_percent", + "margin_fixed_amount", + "margin_percent", + "margin_total_amount", + "total_cost", + ) + }, + ) + span = _conversation_span(["genai", "openinference"], payload) + a = span.attributes + + assert span.dropped_attributes == 0 + assert a[GenAI.REQUEST_MODEL] == "gpt-4o" + assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 + assert a[LiteLLM.TOOLS_DECLARED] == 127 + assert a["llm.input_messages.0.message.content"] == "turn 0" + assert a["llm.output_messages.0.message.content"] == "reply 0" From fcaf2d7d98164fc8561c411a999ffd42469797e9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 10 Sep 2026 08:00:51 +0000 Subject: [PATCH 023/297] fix(otel): size the message ceiling so every vocabulary fits beside it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/mappers/utils.py | 9 ++++++--- .../integrations/otel/test_otel_v2_emitter.py | 15 ++++++++------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index 8d21b774319..d8918491720 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -32,15 +32,18 @@ core telemetry no matter how many vocabularies are configured. """ -MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 4 +MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8 """Span-wide ceiling on attributes spent spelling out chat messages per index. A conversation is the other unbounded family: two attributes per message, for the prompt and the response alike, on the same span. Past a few dozen turns the family alone exceeds the span attribute limit and evicts the core telemetry written before it. The ceiling covers both directions together, since a budget -handed to each direction separately doubles. The complete conversation still -rides the JSON blob attributes; only the per-index convenience keys are capped. +handed to each direction separately doubles. An eighth is the largest share +that still fits beside the tool ceiling and the core of every vocabulary at +once, request parameters, cost breakdown and identity included. The complete +conversation still rides the JSON blob attributes; only the per-index +convenience keys are capped. """ diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index a417bd62124..f571ab7004b 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -461,10 +461,11 @@ def _conversation_payload(turns, choices=1, **overrides): ) -def _conversation_span(mapper_names, payload): +def _conversation_span(mapper_names, payload, legacy_compat=False): """The exported LLM-call span for ``payload`` with content capture on.""" cfg = OpenTelemetryV2Config( exporter="in_memory", + legacy_compat=legacy_compat, mapper_names=list(mapper_names), capture_message_content="span_only", ) @@ -541,13 +542,13 @@ def test_message_cap_is_shared_across_input_and_output(): ) == (MAX_MESSAGE_ATTRS_PER_SPAN // 2) -def test_fully_populated_arize_span_stays_within_the_attribute_limit(): +def test_fully_populated_span_with_every_vocabulary_stays_within_the_attribute_limit(): """Every capped family maxed at once still leaves the whole core intact. - The Arize / Phoenix composition (``genai`` + ``openinference`` + ``legacy``) - with every request parameter, every cost component, a hundred-plus tools, a - two-hundred-turn prompt and twenty choices is the worst case the two - span-wide ceilings have to absorb together. + Every vocabulary in the registry plus ``legacy``, every request parameter, + every cost component, a hundred-plus tools, a two-hundred-turn prompt and + twenty choices is the worst case the two span-wide ceilings have to absorb + together. """ payload = _conversation_payload( 200, @@ -579,7 +580,7 @@ def test_fully_populated_arize_span_stays_within_the_attribute_limit(): ) }, ) - span = _conversation_span(["genai", "openinference"], payload) + span = _conversation_span(["genai", "openinference", "langfuse", "weave", "langtrace"], payload, legacy_compat=True) a = span.attributes assert span.dropped_attributes == 0 From 302a8d43da054f98b7ccb75baf43a8a79bb3ea40 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 10 Sep 2026 22:02:58 +0000 Subject: [PATCH 024/297] fix(cost): bill cached realtime audio tokens at the audio cache-read rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 78 +++++++++++++------ .../litellm_core_utils/llm_cost_calc/utils.py | 40 ++++++++-- ...odel_prices_and_context_window_backup.json | 3 + .../transformation.py | 45 ++++++----- litellm/responses/utils.py | 3 + litellm/types/llms/openai.py | 14 ++++ litellm/types/utils.py | 6 ++ model_prices_and_context_window.json | 3 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 68 ++++++++++++++++ .../responses/test_responses_utils.py | 41 ++++++++++ tests/test_litellm/test_cost_calculator.py | 61 +++++++++++++++ 11 files changed, 314 insertions(+), 48 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..7cd3ea8f303 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -108,6 +108,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.rerank import RerankBilledUnits, RerankResponse from litellm.types.utils import ( + CachedTokensDetails, CallTypesLiteral, LiteLLMRealtimeStreamLoggingObject, LlmProviders, @@ -2310,6 +2311,60 @@ def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str] return [attr for attr in field_names if attr != "cache_creation_tokens"] +def _combine_cached_tokens_details( + current: CachedTokensDetails | None, new: CachedTokensDetails +) -> CachedTokensDetails: + def _sum_optional(current_value: int | None, new_value: int | None) -> int | None: + if current_value is None and new_value is None: + return None + return (current_value or 0) + (new_value or 0) + + return CachedTokensDetails( + text_tokens=_sum_optional( + current.text_tokens if current is not None else None, new.text_tokens + ), + audio_tokens=_sum_optional( + current.audio_tokens if current is not None else None, new.audio_tokens + ), + image_tokens=_sum_optional( + current.image_tokens if current is not None else None, new.image_tokens + ), + ) + + +def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: + if not (hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details): + return + if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: + combined.prompt_tokens_details = PromptTokensDetailsWrapper() + + # Check what keys exist in the model's prompt_tokens_details + # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings + for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): + if ( + hasattr(usage.prompt_tokens_details, attr) + and not attr.startswith("_") + and not callable(_attribute_value(usage.prompt_tokens_details, attr)) + ): + current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 + new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 + if new_val is not None and isinstance(new_val, (int, float)): + setattr( + combined.prompt_tokens_details, + attr, + current_val + new_val, + ) + + new_cached_tokens_details: Final = getattr( + usage.prompt_tokens_details, "cached_tokens_details", None + ) + if isinstance(new_cached_tokens_details, CachedTokensDetails): + combined.prompt_tokens_details.cached_tokens_details = _combine_cached_tokens_details( + getattr(combined.prompt_tokens_details, "cached_tokens_details", None), + new_cached_tokens_details, + ) + + class BaseTokenUsageProcessor: @staticmethod def combine_usage_objects(usage_objects: list[Usage]) -> Usage: @@ -2318,7 +2373,6 @@ class BaseTokenUsageProcessor: """ from litellm.types.utils import ( CompletionTokensDetailsWrapper, - PromptTokensDetailsWrapper, Usage, ) @@ -2337,27 +2391,7 @@ class BaseTokenUsageProcessor: and isinstance(current_val, (int, float)) ): setattr(combined, attr, current_val + new_val) - # Handle nested prompt_tokens_details - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: - combined.prompt_tokens_details = PromptTokensDetailsWrapper() - - # Check what keys exist in the model's prompt_tokens_details - # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings - for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): - if ( - hasattr(usage.prompt_tokens_details, attr) - and not attr.startswith("_") - and not callable(_attribute_value(usage.prompt_tokens_details, attr)) - ): - current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 - new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 - if new_val is not None and isinstance(new_val, (int, float)): - setattr( - combined.prompt_tokens_details, - attr, - current_val + new_val, - ) + _combine_prompt_tokens_details(combined, usage) # Handle nested completion_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e5977ca4156..18ef99597a0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -9,6 +9,8 @@ from types import MappingProxyType from typing import Any, Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError +from typing_extensions import ReadOnly + import litellm from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger @@ -772,6 +774,7 @@ def calculate_cache_writing_cost( class PromptTokensDetailsResult(TypedDict): cache_hit_tokens: int + cache_hit_audio_tokens: ReadOnly[int] cache_creation_tokens: int cache_creation_token_details: CacheCreationTokenDetails | None text_tokens: int @@ -802,12 +805,26 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or None ) - text_tokens: Final = ( - cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None)) - or 0 # default to prompt tokens, if this field is not set + cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) + cached_text_tokens: Final = _get_token_detail_value(cached_tokens_details, "text_tokens") or 0 + cached_audio_tokens: Final = _get_token_detail_value(cached_tokens_details, "audio_tokens") or 0 + cached_image_tokens: Final = _get_token_detail_value(cached_tokens_details, "image_tokens") or 0 + text_tokens: Final = max( + ( + cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None)) + or 0 # default to prompt tokens, if this field is not set + ) + - cached_text_tokens, + 0, + ) + audio_tokens: Final = max( + (cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0) - cached_audio_tokens, + 0, + ) + image_tokens: Final = max( + (cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0) - cached_image_tokens, + 0, ) - audio_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0 - image_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0 video_tokens: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0)) character_count: Final = ( cast( @@ -835,6 +852,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, + cache_hit_audio_tokens=min(cached_audio_tokens, cache_hit_tokens), cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, text_tokens=text_tokens, @@ -918,7 +936,16 @@ def _calculate_input_cost( prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost ### CACHE READ COST - Now uses tiered pricing - prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost + cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"] + audio_cache_read_rate: Final = _get_cost_per_unit( + model_info, + _get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier), + None, + ) + prompt_cost += float(prompt_tokens_details["cache_hit_tokens"] - cache_hit_audio_tokens) * cache_read_cost + prompt_cost += float(cache_hit_audio_tokens) * ( + audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost + ) ### AUDIO COST if prompt_tokens_details["audio_tokens"]: @@ -1149,6 +1176,7 @@ def generic_cost_per_token( ### PROCESSING COST prompt_tokens_details = PromptTokensDetailsResult( cache_hit_tokens=0, + cache_hit_audio_tokens=0, cache_creation_tokens=0, cache_creation_token_details=None, text_tokens=usage.prompt_tokens, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1c0bd32d782..7f4ca991bd3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32502,6 +32502,7 @@ }, "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, @@ -32535,6 +32536,7 @@ }, "gpt-realtime-1.5": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -32702,6 +32704,7 @@ }, "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..cc84c68b0f5 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -45,6 +45,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ) from litellm.types.llms.openai import ( AllMessageValues, + CachedTokensDetails, ChatCompletionImageObject, ChatCompletionImageUrlObject, ChatCompletionRedactedThinkingBlock, @@ -2681,27 +2682,31 @@ class LiteLLMCompletionResponsesConfig: # Translate prompt_tokens_details to input_tokens_details if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: prompt_details: Final = usage.prompt_tokens_details - input_details_dict: Final[dict[str, int]] = {} - - if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None: - input_details_dict["cached_tokens"] = prompt_details.cached_tokens - else: - input_details_dict["cached_tokens"] = 0 - - if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None: - input_details_dict["text_tokens"] = prompt_details.text_tokens - - if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None: - input_details_dict["audio_tokens"] = prompt_details.audio_tokens - - cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) or getattr( - prompt_details, "cache_creation_tokens", None + cached_tokens_details: Final = getattr(prompt_details, "cached_tokens_details", None) + response_usage.input_tokens_details = InputTokensDetails( + cached_tokens=( + prompt_details.cached_tokens + if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None + else 0 + ), + text_tokens=( + prompt_details.text_tokens + if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None + else None + ), + audio_tokens=( + prompt_details.audio_tokens + if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None + else None + ), + cache_write_tokens=( + getattr(prompt_details, "cache_write_tokens", None) + or getattr(prompt_details, "cache_creation_tokens", None) + ), + cached_tokens_details=( + cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None + ), ) - if cache_write_tokens is not None: - input_details_dict["cache_write_tokens"] = cache_write_tokens - - if input_details_dict: - response_usage.input_tokens_details = InputTokensDetails(**input_details_dict) # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 599e978df6a..d63e3ddf0aa 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1179,6 +1179,9 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), + cached_tokens_details=getattr( + response_api_usage.input_tokens_details, "cached_tokens_details", None + ), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b7c4371f32f..274747b4193 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1284,9 +1284,16 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class CachedTokensDetails(BaseModel): + text_tokens: int | None = None + audio_tokens: int | None = None + image_tokens: int | None = None + + class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 + cached_tokens_details: CachedTokensDetails | None = None text_tokens: int | None = None model_config = {"extra": "allow"} @@ -2204,10 +2211,17 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict): transcript: ReadOnly[str] +class OpenAIRealtimeCachedTokensDetails(TypedDict, total=False): + text_tokens: ReadOnly[int] + audio_tokens: ReadOnly[int] + image_tokens: ReadOnly[int] + + class OpenAIRealtimeUsageTokenDetails(TypedDict): audio_tokens: ReadOnly[int] text_tokens: ReadOnly[int] cached_tokens: NotRequired[ReadOnly[int]] + cached_tokens_details: NotRequired[ReadOnly[OpenAIRealtimeCachedTokensDetails]] class OpenAIRealtimeResponseUsage(TypedDict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ab0cc5f959c..39100031dcf 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -58,6 +58,7 @@ from .llms.base import HiddenParams from .llms.openai import ( AllMessageValues, Batch, + CachedTokensDetails, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionRedactedThinkingBlock, @@ -1707,6 +1708,9 @@ class PromptTokensDetailsWrapper( cache_creation_token_details: CacheCreationTokenDetails | None = None """Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching.""" + cached_tokens_details: CachedTokensDetails | None = None + """Details of cached (cache-hit) tokens sent to the model. OpenAI realtime naming; carries the per-modality cache-read split.""" + def __setattr__(self, name: str, value: object) -> None: super().__setattr__(name, value) if name == "cache_write_tokens": @@ -1753,6 +1757,8 @@ class PromptTokensDetailsWrapper( del self.cache_creation_tokens if self.cache_creation_token_details is None: del self.cache_creation_token_details + if self.cached_tokens_details is None: + del self.cached_tokens_details class ServerToolUse(BaseModel): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1c0bd32d782..7f4ca991bd3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32502,6 +32502,7 @@ }, "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, @@ -32535,6 +32536,7 @@ }, "gpt-realtime-1.5": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -32702,6 +32704,7 @@ }, "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index fbb9d178390..3709b526c3b 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2648,6 +2648,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): prompt_tokens_details: PromptTokensDetailsResult = { "cache_hit_tokens": 0, + "cache_hit_audio_tokens": 0, "cache_creation_tokens": 0, "cache_creation_token_details": CacheCreationTokenDetails( ephemeral_5m_input_tokens=100, @@ -5147,3 +5148,70 @@ def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output( assert completion_cost == pytest.approx( 30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] ) + + +def test_cached_realtime_audio_tokens_billed_at_audio_cache_read_rate( + _local_model_cost_map: None, +) -> None: + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, + audio_tokens=167, + cached_tokens=192, + cached_tokens_details={"text_tokens": 64, "audio_tokens": 128}, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" + ) + assert prompt_cost == pytest.approx(0.0015328) + + +def test_prompt_tokens_details_without_cached_tokens_details_unchanged( + _local_model_cost_map: None, +) -> None: + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, audio_tokens=167, cached_tokens=192 + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" + ) + assert prompt_cost == pytest.approx(0.0029888) + + +def test_cached_audio_tokens_fall_back_to_cache_read_input_token_cost() -> None: + model_info: ModelInfo = { + "input_cost_per_token": 4e-6, + "input_cost_per_audio_token": 32e-6, + "cache_read_input_token_cost": 5e-7, + } + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, + audio_tokens=167, + cached_tokens=192, + cached_tokens_details={"text_tokens": 64, "audio_tokens": 128}, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="some-realtime-model", + usage=usage, + custom_llm_provider="openai", + model_info=model_info, + ) + expected = 52 * 4e-6 + 64 * 5e-7 + 39 * 32e-6 + 128 * 5e-7 + assert prompt_cost == pytest.approx(expected) diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 9d9eefdceb3..4d06b5e7bdc 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -577,6 +577,47 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details is not None assert result.completion_tokens_details.reasoning_tokens == 4 + def test_transform_realtime_usage_dict_keeps_cached_tokens_details(self): + usage = { + "input_tokens": 283, + "output_tokens": 0, + "total_tokens": 283, + "input_token_details": { + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens == 192 + assert result.prompt_tokens_details.cached_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens_details.audio_tokens == 128 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 64 + + def test_transform_response_api_usage_object_keeps_cached_tokens_details(self): + usage = ResponseAPIUsage( + input_tokens=283, + output_tokens=0, + total_tokens=283, + input_tokens_details={ + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + ) + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens_details.audio_tokens == 128 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 64 + class TestResponsesAPIProviderSpecificParams: """ diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f610821e06a..0b339e3d525 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4768,3 +4768,64 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> assert combined.completion_tokens_details.reasoning_tokens == 95 assert combined.completion_tokens_details.text_tokens == 38 assert combined.completion_tokens_details.audio_tokens == 0 + + +def test_realtime_combine_sums_nested_cached_tokens_details(): + results: OpenAIRealtimeStreamList = [ + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 283, + "output_tokens": 0, + "total_tokens": 283, + "input_token_details": { + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + } + }, + }, + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 150, + "output_tokens": 0, + "total_tokens": 150, + "input_token_details": { + "text_tokens": 50, + "audio_tokens": 100, + "cached_tokens": 100, + "cached_tokens_details": {"audio_tokens": 100}, + }, + } + }, + }, + ] + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.cached_tokens == 292 + assert combined.prompt_tokens_details.cached_tokens_details is not None + assert combined.prompt_tokens_details.cached_tokens_details.audio_tokens == 228 + assert combined.prompt_tokens_details.cached_tokens_details.text_tokens == 64 + assert combined.prompt_tokens_details.cached_tokens_details.image_tokens is None + + +def test_usage_without_cached_tokens_details_omits_key(): + usage = Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=10), + ) + + dumped = usage.prompt_tokens_details.model_dump() + assert "cached_tokens_details" not in dumped + assert "cached_tokens_details" not in usage.prompt_tokens_details.model_dump_json() From 67fc9e4e3dcb94035c4b4d07d63c3565939d9a3f Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 10 Sep 2026 22:10:06 +0000 Subject: [PATCH 025/297] fix(responses): only emit cache_write_tokens when reported Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 16 +++------- .../transformation.py | 30 +++++++------------ .../test_litellm_completion_responses.py | 2 ++ 3 files changed, 16 insertions(+), 32 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 7cd3ea8f303..8daa2de416b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2320,15 +2320,9 @@ def _combine_cached_tokens_details( return (current_value or 0) + (new_value or 0) return CachedTokensDetails( - text_tokens=_sum_optional( - current.text_tokens if current is not None else None, new.text_tokens - ), - audio_tokens=_sum_optional( - current.audio_tokens if current is not None else None, new.audio_tokens - ), - image_tokens=_sum_optional( - current.image_tokens if current is not None else None, new.image_tokens - ), + text_tokens=_sum_optional(current.text_tokens if current is not None else None, new.text_tokens), + audio_tokens=_sum_optional(current.audio_tokens if current is not None else None, new.audio_tokens), + image_tokens=_sum_optional(current.image_tokens if current is not None else None, new.image_tokens), ) @@ -2355,9 +2349,7 @@ def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: current_val + new_val, ) - new_cached_tokens_details: Final = getattr( - usage.prompt_tokens_details, "cached_tokens_details", None - ) + new_cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) if isinstance(new_cached_tokens_details, CachedTokensDetails): combined.prompt_tokens_details.cached_tokens_details = _combine_cached_tokens_details( getattr(combined.prompt_tokens_details, "cached_tokens_details", None), diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index cc84c68b0f5..e6f90b99b60 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2683,30 +2683,20 @@ class LiteLLMCompletionResponsesConfig: if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: prompt_details: Final = usage.prompt_tokens_details cached_tokens_details: Final = getattr(prompt_details, "cached_tokens_details", None) - response_usage.input_tokens_details = InputTokensDetails( - cached_tokens=( - prompt_details.cached_tokens - if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None - else 0 - ), - text_tokens=( - prompt_details.text_tokens - if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None - else None - ), - audio_tokens=( - prompt_details.audio_tokens - if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None - else None - ), - cache_write_tokens=( - getattr(prompt_details, "cache_write_tokens", None) - or getattr(prompt_details, "cache_creation_tokens", None) - ), + cache_write_tokens: Final = getattr(prompt_details, "cache_write_tokens", None) or getattr( + prompt_details, "cache_creation_tokens", None + ) + input_tokens_details: Final = InputTokensDetails( + cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, + text_tokens=prompt_details.text_tokens, + audio_tokens=prompt_details.audio_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), ) + if cache_write_tokens is not None: + setattr(input_tokens_details, "cache_write_tokens", cache_write_tokens) + response_usage.input_tokens_details = input_tokens_details # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 46249e50572..be96c2a4bf5 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2605,6 +2605,7 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is not None assert response_usage.input_tokens_details.cached_tokens == 5 assert response_usage.input_tokens_details.text_tokens == 8 + assert "cache_write_tokens" not in response_usage.input_tokens_details.model_dump() def test_transform_usage_with_cached_tokens_gemini(self): """Test that cached_tokens from Gemini are properly transformed to input_tokens_details""" @@ -2667,6 +2668,7 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is not None assert response_usage.input_tokens_details.cached_tokens == 100 assert getattr(response_usage.input_tokens_details, "cache_write_tokens", None) == 800 + assert response_usage.input_tokens_details.model_dump()["cache_write_tokens"] == 800 def test_transform_usage_with_reasoning_tokens_gemini(self): """Test that reasoning_tokens from Gemini are properly transformed to output_tokens_details""" From 5737cab258b405689a55e1e6dcaec385673fa572 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 10 Sep 2026 22:27:58 +0000 Subject: [PATCH 026/297] fix(cost): cap nested cached modality counts at cached_tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 -- .../litellm_core_utils/llm_cost_calc/utils.py | 16 +++++++++++---- .../llm_cost_calc/test_llm_cost_calc_utils.py | 20 +++++++++++++++++++ 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8daa2de416b..b865318f3af 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2332,8 +2332,6 @@ def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: combined.prompt_tokens_details = PromptTokensDetailsWrapper() - # Check what keys exist in the model's prompt_tokens_details - # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): if ( hasattr(usage.prompt_tokens_details, attr) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 18ef99597a0..dc689ca9618 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -806,9 +806,17 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: or None ) cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) - cached_text_tokens: Final = _get_token_detail_value(cached_tokens_details, "text_tokens") or 0 - cached_audio_tokens: Final = _get_token_detail_value(cached_tokens_details, "audio_tokens") or 0 - cached_image_tokens: Final = _get_token_detail_value(cached_tokens_details, "image_tokens") or 0 + cached_audio_tokens: Final = min( + _get_token_detail_value(cached_tokens_details, "audio_tokens") or 0, cache_hit_tokens + ) + cached_text_tokens: Final = min( + _get_token_detail_value(cached_tokens_details, "text_tokens") or 0, + cache_hit_tokens - cached_audio_tokens, + ) + cached_image_tokens: Final = min( + _get_token_detail_value(cached_tokens_details, "image_tokens") or 0, + cache_hit_tokens - cached_audio_tokens - cached_text_tokens, + ) text_tokens: Final = max( ( cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None)) @@ -852,7 +860,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, - cache_hit_audio_tokens=min(cached_audio_tokens, cache_hit_tokens), + cache_hit_audio_tokens=cached_audio_tokens, cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, text_tokens=text_tokens, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3709b526c3b..33825c8dd01 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5215,3 +5215,23 @@ def test_cached_audio_tokens_fall_back_to_cache_read_input_token_cost() -> None: ) expected = 52 * 4e-6 + 64 * 5e-7 + 39 * 32e-6 + 128 * 5e-7 assert prompt_cost == pytest.approx(expected) + + +def test_cached_audio_tokens_capped_at_cached_tokens(_local_model_cost_map: None) -> None: + """Nested cached_tokens_details exceeding cached_tokens must not over-subtract the audio bucket.""" + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, + audio_tokens=167, + cached_tokens=100, + cached_tokens_details={"audio_tokens": 128}, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" + ) + assert prompt_cost == pytest.approx(116 * 4e-6 + (167 - 100) * 32e-6 + 100 * 4e-7) From e205be80df1e23a22b4a5eee5f37bc9d52cb1ce1 Mon Sep 17 00:00:00 2001 From: oliver Date: Fri, 11 Sep 2026 07:41:13 +0000 Subject: [PATCH 027/297] fix(proxy): enforce custom_key_update policy on /key/regenerate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 47 ++++-- .../test_key_management_endpoints.py | 150 ++++++++++++++++++ 2 files changed, 184 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 749a940de0e..106e9b8027b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -294,6 +294,35 @@ def _custom_key_update_hook( return hooks.user_custom_key_update +async def _enforce_custom_key_update_policy( + hook: Callable[..., Awaitable[Mapping[str, object]]] | None, + data: UpdateKeyRequest, +) -> None: + if hook is None: + return + if not inspect.iscoroutinefunction(hook): + raise ValueError("user_custom_key_update must be a coroutine") + result: Final = await hook(data) + if not result.get("decision", True): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=result.get("message", "Authentication Failed - Custom Auth Rule"), + ) + + +def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: + changed_fields: Final = MappingProxyType( + { + field: value + for field, value in data.model_dump(exclude_unset=True).items() + if field in UpdateKeyRequest.model_fields and field != "key" + } + ) + if not changed_fields: + return None + return UpdateKeyRequest(key=key, **changed_fields) + + class _LegacyDumpable(Protocol): def dict(self) -> Mapping[str, object]: ... @@ -3066,19 +3095,7 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Custom key update hook - custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook( - proxy_server - ) - if custom_key_update_hook is not None: - if inspect.iscoroutinefunction(custom_key_update_hook): - result: Final = await custom_key_update_hook(data) - else: - raise ValueError("user_custom_key_update must be a coroutine") - decision: Final = result.get("decision", True) - message: Final = result.get("message", "Authentication Failed - Custom Auth Rule") - if not decision: - raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) + await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=data) # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) @@ -5069,6 +5086,7 @@ async def _execute_virtual_key_regeneration( proxy_logging_obj: ProxyLogging, ) -> GenerateKeyResponse: """Generate new token, update DB, invalidate cache, and return response.""" + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import hash_token # Mirror the /key/update ownership rebind guard. See helper docstring. @@ -5116,6 +5134,9 @@ async def _execute_virtual_key_regeneration( non_default_values = {} if data is not None: + update_request: Final = _regenerate_request_as_update_request(key=hashed_api_key, data=data) + if update_request is not None: + await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request) # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 65cc23ea67f..fbe790dbed5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -24,18 +24,21 @@ from litellm.proxy._types import ( LitellmUserRoles, Member, ProxyException, + RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, _check_project_key_limits, _check_team_key_limits, _common_key_generation_helper, _enforce_upperbound_key_params, + _execute_virtual_key_regeneration, _get_and_validate_existing_key, _list_key_helper, _persist_deleted_verification_tokens, @@ -12011,6 +12014,153 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_hook_denies(): + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="3000d") + mock_prisma_client = _make_regenerate_mock_prisma() + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ) as insert_deprecated_key, + patch( # test-quality-ok: cache eviction is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "duration must be <= 7d" + insert_deprecated_key.assert_not_awaited() + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert len(received_data) == 1 + assert received_data[0].key == "abc123" + assert received_data[0].duration == "3000d" + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_hook_approves(): + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="5d") + mock_prisma_client = _make_regenerate_mock_prisma() + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for policy approval + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + assert len(received_data) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data): + mock_prisma_client = _make_regenerate_mock_prisma() + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + raise AssertionError(f"custom key update hook called with {data}") + + with ( + patch( # test-quality-ok: deterministic token setup for unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ From 560e5da320bc00636057b3de3edd05d6a2da1745 Mon Sep 17 00:00:00 2001 From: oliver Date: Fri, 11 Sep 2026 08:02:18 +0000 Subject: [PATCH 028/297] fix(proxy): archive the old key only after regenerate validation passes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 18 +++++++---------- .../test_key_management_endpoints.py | 20 +++++++++++++++---- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 106e9b8027b..97e2a4f3be5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5155,6 +5155,13 @@ async def _execute_virtual_key_regeneration( prisma_client=prisma_client, ) + await _persist_deleted_verification_tokens( + keys=[key_in_db], + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( prisma_client=prisma_client, @@ -5453,17 +5460,6 @@ async def regenerate_key_fn( if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): litellm_changed_by = None - # Save the old key record to deleted table before regeneration. - # This preserves key_alias and team_id metadata for historical spend records. - # If this fails, abort the regeneration to avoid permanently losing the - # old hash→metadata mapping. - await _persist_deleted_verification_tokens( - keys=[_key_in_db], - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - return await _execute_virtual_key_regeneration( prisma_client=prisma_client, key_in_db=_key_in_db, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index fbe790dbed5..6ec04815420 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -11935,6 +11935,10 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ), + patch( # test-quality-ok: archival path is outside upperbound rejection + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -11955,6 +11959,7 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk assert exc_info.value.status_code == 400 assert "duration" in str(exc_info.value.detail) # Rejected regenerate must not reach the DB update. + persist_deleted_verification_tokens.assert_not_awaited() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 @@ -12037,6 +12042,10 @@ async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_h "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ) as insert_deprecated_key, + patch( # test-quality-ok: archival path is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( # test-quality-ok: cache eviction is outside policy rejection "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -12063,6 +12072,7 @@ async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_h assert exc_info.value.status_code == 403 assert exc_info.value.detail == "duration must be <= 7d" insert_deprecated_key.assert_not_awaited() + persist_deleted_verification_tokens.assert_not_awaited() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 assert len(received_data) == 1 assert received_data[0].key == "abc123" @@ -12092,6 +12102,10 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ), + patch( # test-quality-ok: verify archival follows policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( # test-quality-ok: cache eviction is outside policy approval "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -12115,6 +12129,8 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho ) assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + persist_deleted_verification_tokens.assert_awaited_once() + assert persist_deleted_verification_tokens.call_args.kwargs["keys"] == [existing_key] assert len(received_data) == 1 @@ -13922,10 +13938,6 @@ async def test_regenerate_applies_normalized_mcp_object_permission(): "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_vector_stores_against_team", new_callable=AsyncMock, ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", - new_callable=AsyncMock, - ), patch( "litellm.proxy.management_endpoints.key_management_endpoints._execute_virtual_key_regeneration", execute_mock, From 135ec00b27ac2452611032336c743c57f64ccae7 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:06:09 +0000 Subject: [PATCH 029/297] feat(model_armor): logging_only mode scans completed streams after delivery Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 25 +- .../model_armor/model_armor.py | 68 ++++- .../integrations/test_custom_guardrail.py | 57 +++- .../guardrail_hooks/test_model_armor.py | 289 ++++++++++++++++++ 4 files changed, 428 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 77bf4820a1a..ae6d646b5ae 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -30,6 +30,7 @@ from litellm.types.utils import ( GuardrailStatus, GuardrailTracingDetail, LLMResponseTypes, + ModelResponse, StandardLoggingGuardrailInformation, ) @@ -602,9 +603,7 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) - if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks - else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) if self.uses_apply_guardrail_interface() else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( @@ -883,7 +882,9 @@ class CustomGuardrail(CustomLogger): """logging_only: run apply_guardrail on copies of the logged request/response and record the verdict.""" from litellm.llms import get_guardrail_translation_mapping - if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: + if not self.uses_apply_guardrail_interface(): + return kwargs, result + if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only): return kwargs, result try: translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))() @@ -922,6 +923,8 @@ class CustomGuardrail(CustomLogger): translation: "BaseTranslation", scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: + from litellm.llms import get_guardrail_translation_mapping + optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -933,8 +936,18 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - await translation.process_output_response( - response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request + response: Final = ( + kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result + ) + if response is None: + return + output_translation: Final = ( + get_guardrail_translation_mapping(CallTypes.acompletion)() + if isinstance(response, ModelResponse) + else translation + ) + await output_translation.process_output_response( + response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request ) def supports_scan_only_tool_results(self) -> bool: diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index fde40111d49..01b05227a91 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,11 +1,13 @@ +import time from collections.abc import AsyncGenerator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional import httpx from fastapi import HTTPException if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel import json @@ -52,6 +54,7 @@ from litellm.types.utils import ( CallTypes, CallTypesLiteral, Choices, + GenericGuardrailAPIInputs, GuardrailStatus, ModelResponse, ModelResponseStream, @@ -118,8 +121,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Supports: - Pre-call sanitization (sanitizeUserPrompt) - Post-call sanitization (sanitizeModelResponse) + - logging_only: scans the completed response after it reaches the client and + records the verdict in spend logs without blocking """ + use_native_lifecycle_hooks: ClassVar[bool] = True + @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ @@ -128,6 +135,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): GuardrailEventHooks.post_call, GuardrailEventHooks.pre_mcp_call, GuardrailEventHooks.during_mcp_call, + GuardrailEventHooks.logging_only, ] def __init__( @@ -1096,6 +1104,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): add_guardrail_to_applied_guardrails_header, ) + if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: + async for chunk in response: + yield chunk + return + all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response]) if not all_chunks or self._is_terminal_error_stream(all_chunks): @@ -1213,6 +1226,59 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): for chunk in all_chunks: yield chunk + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + content: Final = "\n".join(text for text in inputs.get("texts") or () if text) + if not content: + return inputs + + source: Final[Literal["user_prompt", "model_response"]] = ( + "user_prompt" if input_type == "request" else "model_response" + ) + start_time: Final = time.time() + try: + armor_response: Final = await self.make_model_armor_request( + content=content, source=source, request_data=request_data + ) + except ModelArmorAPIError as e: + error_end_time: Final = time.time() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(e), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + guardrail_provider="model_armor", + start_time=start_time, + end_time=error_end_time, + duration=error_end_time - start_time, + ) + raise + + flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) + end_time: Final = time.time() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=self._build_logging_response(armor_response), + request_data=request_data, + guardrail_status="guardrail_flagged" if flagged else "success", + guardrail_provider="model_armor", + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + if flagged: + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail( + "Response blocked by Model Armor" if input_type == "response" else "Violated content safety policy", + armor_response, + ), + ) + return inputs + @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: """ diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index ddc8439a83a..c0463210445 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2404,7 +2404,7 @@ def test_logging_only_requires_framework_support_or_explicit_declaration( event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, ) -> None: supported: Final = [GuardrailEventHooks.pre_call] - if guardrail_type is _InheritedApplyGuardrail: + if guardrail_type is not CustomGuardrail: guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) assert guardrail.event_hook == event_hook assert supported == [GuardrailEventHooks.pre_call] @@ -2566,7 +2566,7 @@ class TestLoggingOnlyApplyGuardrail: assert [e["guardrail_status"] for e in entries] == ["success"] @pytest.mark.asyncio - async def test_native_lifecycle_hook_guardrail_is_left_alone(self): + async def test_native_lifecycle_hook_guardrail_scans_in_logging_only(self): class _NativeHooks(_ApplyOnlyObserver): use_native_lifecycle_hooks = True @@ -2575,9 +2575,9 @@ class TestLoggingOnlyApplyGuardrail: out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) - assert guardrail.calls == [] - assert out_kwargs is kwargs + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] assert out_response is response + assert out_kwargs["standard_logging_object"]["guardrail_information"] @pytest.mark.asyncio async def test_aresponses_scans_logged_messages_when_input_is_cleared(self): @@ -2829,3 +2829,52 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: assert response.choices[0].message.content == "filtered response" assert "guardrail_to_apply" not in request_data assert len(_guardrail_entries(request_data)) == 1 + + +class _NativeLifecycleLoggingGuardrail(CustomGuardrail): + """Native lifecycle guardrail that also implements apply_guardrail, like the azure guards.""" + + use_native_lifecycle_hooks: ClassVar[bool] = True + + def __init__(self): + from litellm.types.guardrails import GuardrailEventHooks + + super().__init__( + guardrail_name="native-logging-guardrail", + event_hook=GuardrailEventHooks.logging_only, + ) + self.calls: list = [] + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.calls.append((input_type, list(inputs.get("texts") or []))) + return inputs + + +@pytest.mark.asyncio +async def test_native_lifecycle_guardrail_logging_only_scans_assembled_response(): + """A use_native_lifecycle_hooks guardrail accepts mode logging_only and its + async_logging_hook scans kwargs["async_complete_streaming_response"], not the raw result.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _NativeLifecycleLoggingGuardrail() + assembled = ModelResponse( + choices=[Choices(message=Message(role="assistant", content="assembled stream text"))] + ) + sentinel_result = object() + kwargs = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + "async_complete_streaming_response": assembled, + } + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=sentinel_result, call_type=CallTypes.acompletion.value + ) + + assert out_result is sentinel_result + assert ("response", ["assembled stream text"]) in guardrail.calls + assert out_kwargs["standard_logging_object"]["guardrail_information"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 47089b7b1b1..a6ed4e14616 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -4929,3 +4929,292 @@ def test_every_responses_delta_event_is_in_the_scanned_set(): } assert not missing assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES + + +def _clean_armor_response() -> dict: + return { + "sanitizationResult": { + "filterMatchState": "NO_MATCH_FOUND", + "filterResults": {}, + } + } + + +def _flagged_armor_response() -> dict: + return { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}, + } + } + + +def _logging_only_guardrail() -> ModelArmorGuardrail: + return ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-logging", + event_hook=GuardrailEventHooks.logging_only, + ) + + +def _logged_kwargs() -> dict: + return { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + + +def _chat_response(text: str) -> litellm.ModelResponse: + return litellm.ModelResponse( + choices=[ + litellm.types.utils.Choices( + message=litellm.types.utils.Message(role="assistant", content=text) + ) + ] + ) + + +def _stream_chunk(text: str) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content=text) + ) + ] + ) + + +def _metadata_entries(kwargs: dict) -> list: + return kwargs["standard_logging_object"].get("guardrail_information") or [] + + +def test_logging_only_mode_is_accepted_and_keeps_native_hooks(): + guardrail = _logging_only_guardrail() + assert guardrail.event_hook == GuardrailEventHooks.logging_only + assert guardrail.use_native_lifecycle_hooks is True + assert GuardrailEventHooks.logging_only in ModelArmorGuardrail.get_supported_event_hooks() + + post_call_guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-post", + event_hook=GuardrailEventHooks.post_call, + ) + assert post_call_guardrail._deployment_hook_target() is post_call_guardrail + + +@pytest.mark.asyncio +async def test_logging_only_stream_yields_chunks_without_waiting_for_scan(): + """A logging_only guardrail must pass stream chunks straight through; the scan happens + afterwards on the assembled response via async_logging_hook.""" + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock( + side_effect=AssertionError("logging_only must not scan the stream") + ) + + produced = 0 + + async def gen(): + nonlocal produced + for i in range(3): + produced += 1 + yield _stream_chunk(f"chunk-{i} ") + + hook_iter = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=gen(), + request_data={"metadata": {}, "guardrails": ["model-armor-logging"]}, + ) + first = await hook_iter.__anext__() + assert produced == 1 + chunks = [first] + async for chunk in hook_iter: + chunks.append(chunk) + assert len(chunks) == 3 + guardrail.make_model_armor_request.assert_not_awaited() + + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + response = _chat_response("all clear") + kwargs = _logged_kwargs() + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + assert out_result is response + entries = _metadata_entries(out_kwargs) + assert len(entries) >= 1 + entry = entries[-1] + assert entry["guardrail_status"] == "success" + assert entry["guardrail_mode"] == "logging_only" + assert entry["guardrail_provider"] == "model_armor" + + +@pytest.mark.asyncio +async def test_logging_only_records_flagged_verdict_without_altering_response(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + response = _chat_response("flagged output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_result is response + entries = _metadata_entries(out_kwargs) + assert entries[-1]["guardrail_status"] == "guardrail_flagged" + assert entries[-1]["guardrail_mode"] == "logging_only" + + +@pytest.mark.asyncio +async def test_logging_only_records_model_armor_api_error(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock( + side_effect=ModelArmorAPIError("Model Armor API error (upstream 500)") + ) + response = _chat_response("some output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_result is response + entries = _metadata_entries(out_kwargs) + assert entries[-1]["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_logging_only_scans_assembled_responses_api_stream(): + """The terminal ResponseCompletedEvent is an envelope; the scan must run on the + assembled ResponsesAPIResponse kept in kwargs.""" + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + assembled = ResponsesAPIResponse( + id="resp-1", + created_at=1700000000, + output=[ + ResponseOutputMessage( + id="msg-1", + type="message", + role="assistant", + status="completed", + content=[ + ResponseOutputText( + annotations=[], text="assembled output text", type="output_text" + ) + ], + ) + ], + ) + event = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=assembled + ) + + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + kwargs = _logged_kwargs() + del kwargs["messages"] + kwargs["input"] = "hello" + kwargs["async_complete_streaming_response"] = assembled + + out_kwargs, _ = await guardrail.async_logging_hook( + kwargs=kwargs, result=event, call_type="aresponses" + ) + + response_scans = [ + call + for call in guardrail.make_model_armor_request.await_args_list + if call.kwargs.get("source") == "model_response" + ] + assert response_scans, "expected a model_response scan of the assembled response" + assert "assembled output text" in response_scans[0].kwargs["content"] + assert _metadata_entries(out_kwargs) + + +@pytest.mark.asyncio +async def test_logging_only_scans_anthropic_messages_model_response(): + """/v1/messages logs a ModelResponse; the output scan must extract the assistant text.""" + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + kwargs = _logged_kwargs() + kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + response = _chat_response("anthropic assembled text") + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="anthropic_messages" + ) + + assert out_result is response + response_scans = [ + call + for call in guardrail.make_model_armor_request.await_args_list + if call.kwargs.get("source") == "model_response" + ] + assert response_scans + assert "anthropic assembled text" in response_scans[0].kwargs["content"] + assert _metadata_entries(out_kwargs) + + +@pytest.mark.asyncio +async def test_logging_only_skips_output_scan_when_no_assembled_response(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + kwargs = _logged_kwargs() + + await guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion") + + sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list] + assert "model_response" not in sources + + +@pytest.mark.asyncio +async def test_native_post_call_mode_ignores_logging_hook(): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-post", + event_hook=GuardrailEventHooks.post_call, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + response = _chat_response("some output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_kwargs is kwargs + assert out_result is response + guardrail.make_model_armor_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_apply_guardrail_raises_on_flagged_content(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + request_data = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["forbidden output"]}, + request_data=request_data, + input_type="response", + ) + + assert exc_info.value.status_code == 400 + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert entries[-1]["guardrail_status"] == "guardrail_flagged" From a5cc65f1a9b5f7a674d98f2e45460790952724e5 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:10:35 +0000 Subject: [PATCH 030/297] refactor(custom_guardrail): resolve logging_only output translation in async_logging_hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index ae6d646b5ae..305a5c20764 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -901,8 +901,16 @@ class CustomGuardrail(CustomLogger): for key, value in (litellm_params.get("metadata") or {}).items() if key != "standard_logging_guardrail_information" } + response: Final = ( + kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result + ) + output_translation: Final = ( + get_guardrail_translation_mapping(CallTypes.acompletion)() + if isinstance(response, ModelResponse) + else translation + ) try: - await self._scan_logged_call(kwargs, result, translation, scratch_metadata) + await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata) except Exception as e: verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) recorded: Final = scratch_metadata.get("standard_logging_guardrail_information") @@ -919,12 +927,11 @@ class CustomGuardrail(CustomLogger): async def _scan_logged_call( self, kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract - result: object, + response: object | None, translation: "BaseTranslation", + output_translation: "BaseTranslation", scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: - from litellm.llms import get_guardrail_translation_mapping - optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -936,16 +943,8 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - response: Final = ( - kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result - ) if response is None: return - output_translation: Final = ( - get_guardrail_translation_mapping(CallTypes.acompletion)() - if isinstance(response, ModelResponse) - else translation - ) await output_translation.process_output_response( response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request ) From f417d7f739fca2b8abed6ad3d4a399e39b2cfee7 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:49:11 +0000 Subject: [PATCH 031/297] fix(model_armor): record logging_only verdicts without raising so both scans run Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_armor/model_armor.py | 12 +---- .../guardrail_hooks/test_model_armor.py | 53 ++++++++++++++++--- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 01b05227a91..127a7ea6786 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1245,7 +1245,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): armor_response: Final = await self.make_model_armor_request( content=content, source=source, request_data=request_data ) - except ModelArmorAPIError as e: + except (ModelArmorAPIError, httpx.HTTPError) as e: error_end_time: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=str(e), @@ -1256,7 +1256,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=error_end_time, duration=error_end_time - start_time, ) - raise + return inputs flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) end_time: Final = time.time() @@ -1269,14 +1269,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) - if flagged: - raise HTTPException( - status_code=400, - detail=self._build_block_error_detail( - "Response blocked by Model Armor" if input_type == "response" else "Violated content safety policy", - armor_response, - ), - ) return inputs @staticmethod diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index a6ed4e14616..bd358e84148 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5203,18 +5203,55 @@ async def test_native_post_call_mode_ignores_logging_hook(): @pytest.mark.asyncio -async def test_apply_guardrail_raises_on_flagged_content(): +async def test_apply_guardrail_records_flagged_without_raising(): guardrail = _logging_only_guardrail() guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) request_data = {"metadata": {}} + inputs = {"texts": ["forbidden output"]} - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs={"texts": ["forbidden output"]}, - request_data=request_data, - input_type="response", - ) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) - assert exc_info.value.status_code == 400 + assert result == inputs entries = request_data["metadata"]["standard_logging_guardrail_information"] assert entries[-1]["guardrail_status"] == "guardrail_flagged" + + +@pytest.mark.asyncio +async def test_logging_only_records_transport_error(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(side_effect=httpx.ConnectError("boom")) + response = _chat_response("some output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_result is response + entries = _metadata_entries(out_kwargs) + failed = [e for e in entries if e["guardrail_status"] == "guardrail_failed_to_respond"] + assert failed + assert all(e["guardrail_provider"] == "model_armor" for e in failed) + + +@pytest.mark.asyncio +async def test_logging_only_flagged_prompt_still_scans_response(): + """A flagged input scan must not abort the output scan; both verdicts are recorded.""" + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + response = _chat_response("flagged output") + kwargs = _logged_kwargs() + + out_kwargs, _ = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list] + assert sources == ["user_prompt", "model_response"] + entries = _metadata_entries(out_kwargs) + flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] + assert len(flagged) == 2 From 39b916f13f3330b5099ff45ede27e59f7278f53a Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 09:07:06 +0000 Subject: [PATCH 032/297] fix(model_armor): decorate apply_guardrail with log_guardrail_information Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/guardrails/guardrail_hooks/model_armor/model_armor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 127a7ea6786..a833de6096d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1226,6 +1226,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): for chunk in all_chunks: yield chunk + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, From 2ca29a9a9157d52149a4c2cd38e98f89c7f57e30 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 09:33:57 +0000 Subject: [PATCH 033/297] fix(model_armor): gate apply_guardrail raise to non-logging_only and require native guardrails to declare logging_only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 4 ++- .../model_armor/model_armor.py | 12 +++++++-- .../integrations/test_custom_guardrail.py | 2 +- .../guardrail_hooks/test_model_armor.py | 27 +++++++++++++++++++ 4 files changed, 41 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 305a5c20764..407445b2828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -603,7 +603,9 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) if self.uses_apply_guardrail_interface() else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index a833de6096d..a0563a7a1c9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,7 @@ import time from collections.abc import AsyncGenerator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal import httpx from fastapi import HTTPException @@ -1232,7 +1232,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional["LiteLLMLoggingObj"] = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: content: Final = "\n".join(text for text in inputs.get("texts") or () if text) if not content: @@ -1270,6 +1270,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) + if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only): + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail( + "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor", + armor_response, + ), + ) return inputs @staticmethod diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index c0463210445..bdb7fad21b3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2404,7 +2404,7 @@ def test_logging_only_requires_framework_support_or_explicit_declaration( event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, ) -> None: supported: Final = [GuardrailEventHooks.pre_call] - if guardrail_type is not CustomGuardrail: + if guardrail_type is _InheritedApplyGuardrail: guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) assert guardrail.event_hook == event_hook assert supported == [GuardrailEventHooks.pre_call] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index bd358e84148..126e162fec8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5255,3 +5255,30 @@ async def test_logging_only_flagged_prompt_still_scans_response(): entries = _metadata_entries(out_kwargs) flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] assert len(flagged) == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): + """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a + non-logging_only instance must signal the block so flagged text is not returned as clean.""" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-pre", + event_hook=GuardrailEventHooks.pre_call, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + request_data = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["forbidden prompt"]}, + request_data=request_data, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + entries = request_data["metadata"]["standard_logging_guardrail_information"] + flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] + assert len(flagged) == 1 From a067557dae5c0bb52fdb11176437a39c9a0d9ac3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 11 Sep 2026 23:16:53 +0000 Subject: [PATCH 034/297] fix(otel): index the opener and the latest prompt turns, not the oldest A value length limit clips the input.value blob, so the per-index keys are the only untruncated copy of a message. Indexing the leading prompt messages dropped the live user turn from every span attribute on long conversations. Keep message 0 and the most recent turns under the same span-wide budget, original indices preserved, reply reservation unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../otel/mappers/openinference.py | 31 +++++++++++++------ .../integrations/otel/test_otel_v2_emitter.py | 29 ++++++++++++++++- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index 1e2dbf6974d..1fa19b8d4a5 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -92,8 +92,13 @@ class OpenInferenceMapper: return { **collect(self._LLM_CALL_ATTRS, data), **collect(self._BLOB_ATTRS, data), - **self._messages("llm.input_messages", "input.value", data.messages_in, indexed_in), - **self._messages("llm.output_messages", "output.value", outputs, indexed_out), + **self._messages( + "llm.input_messages", + "input.value", + data.messages_in, + self._prompt_positions(len(data.messages_in), indexed_in), + ), + **self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)), **self._tools(data), } @@ -109,18 +114,26 @@ class OpenInferenceMapper: return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out @staticmethod - def _messages(prefix: str, value_key: str, messages: Sequence[object], indexed: int) -> AttributeMap: - """``{prefix}.{idx}.message.*`` keys for the leading ``indexed`` messages + the ``value_key`` blob of all.""" + def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]: + """Which prompt messages get per-index attributes: message 0 and the most recent turns. + + A value length limit clips the ``input.value`` blob, so the system prompt and the + live turn each keep a short key of their own. The middle of a long prompt does not. + """ + if total <= indexed: + return tuple(range(total)) + return (0, *range(total - indexed + 1, total)) + + @staticmethod + def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap: + """``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all.""" parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages] attrs: Final = drop_none( { key: value - for idx, (role, content) in enumerate(parsed[:indexed]) + for idx, (role, content) in ((idx, parsed[idx]) for idx in positions) for key, value in ( - ( - f"{prefix}.{idx}.message.role", - role if isinstance(role, str) else None, - ), + (f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None), (f"{prefix}.{idx}.message.content", content), ) } diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index f571ab7004b..6491ab0f79f 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -505,7 +505,8 @@ def test_long_conversation_does_not_evict_core_attributes(turns): assert a["llm.input_messages.0.message.content"] == "turn 0" assert a["llm.output_messages.0.message.content"] == "reply 0" - assert f"llm.input_messages.{turns - 1}.message.role" not in a + assert a[f"llm.input_messages.{turns - 1}.message.content"] == f"turn {turns - 1}" + assert f"llm.input_messages.{turns // 2}.message.role" not in a assert len(json.loads(a["input.value"])) == turns assert len(json.loads(a["output.value"])) == 1 assert len(json.loads(a[GenAI.INPUT_MESSAGES])) == turns @@ -520,6 +521,31 @@ def test_short_conversation_keeps_every_message_indexed(): assert a[f"llm.output_messages.{idx}.message.content"] == f"reply {idx}" +def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit(monkeypatch): + """The per-index keys are the only untruncated copy once the SDK clips string values. + + Operators bound attribute sizes with ``OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT``, + which cuts the ``input.value`` blob short. The system prompt and the live turn + then have to survive as their own short keys, whatever the conversation length. + """ + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") + payload = _conversation_payload(60) + payload["messages"][0] = {"role": "system", "content": "be terse"} + payload["messages"][-1] = {"role": "user", "content": "LATEST-TURN"} + a = _conversation_span(["genai", "openinference"], payload).attributes + + assert len(a["input.value"]) == 256 + assert a["llm.input_messages.0.message.role"] == "system" + assert a["llm.input_messages.0.message.content"] == "be terse" + assert a["llm.input_messages.59.message.role"] == "user" + assert a["llm.input_messages.59.message.content"] == "LATEST-TURN" + assert a["llm.output_messages.0.message.content"] == "reply 0" + assert [int(key.split(".")[2]) for key in a if key.endswith("message.content") and key.startswith("llm.input_")] == [ + 0, + *range(54, 60), + ] + + def test_message_cap_is_shared_across_input_and_output(): """One span-wide allowance covers both directions, and the response always keeps a share. @@ -588,4 +614,5 @@ def test_fully_populated_span_with_every_vocabulary_stays_within_the_attribute_l assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 assert a[LiteLLM.TOOLS_DECLARED] == 127 assert a["llm.input_messages.0.message.content"] == "turn 0" + assert a["llm.input_messages.199.message.content"] == "turn 199" assert a["llm.output_messages.0.message.content"] == "reply 0" From a15309dfe820836a41e914228359d7b5becc3744 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 11 Sep 2026 23:32:44 +0000 Subject: [PATCH 035/297] refactor(otel): trim the message cap docstrings to one line each Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../otel/mappers/openinference.py | 13 ++------- litellm/integrations/otel/mappers/utils.py | 13 ++------- .../integrations/otel/test_otel_v2_emitter.py | 29 +++---------------- 3 files changed, 9 insertions(+), 46 deletions(-) diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index 1fa19b8d4a5..a7e0f1af3ac 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -104,22 +104,13 @@ class OpenInferenceMapper: @staticmethod def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]: - """How many prompt and response messages get per-index attributes. - - Both directions share one span-wide allowance. The response is reserved at - least half of it, so a long prompt can never push the completion off the - span, and the prompt takes whatever the response leaves unused. - """ + """Prompt and response share one allowance; the response is reserved at least half of it.""" indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs)) return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out @staticmethod def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]: - """Which prompt messages get per-index attributes: message 0 and the most recent turns. - - A value length limit clips the ``input.value`` blob, so the system prompt and the - live turn each keep a short key of their own. The middle of a long prompt does not. - """ + """Prompt messages that get per-index attributes: message 0 and the most recent turns.""" if total <= indexed: return tuple(range(total)) return (0, *range(total - indexed + 1, total)) diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index d8918491720..c023621d2ef 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -33,17 +33,10 @@ core telemetry no matter how many vocabularies are configured. MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8 -"""Span-wide ceiling on attributes spent spelling out chat messages per index. +"""Span-wide ceiling on per-index chat message attributes, prompt and response together. -A conversation is the other unbounded family: two attributes per message, for -the prompt and the response alike, on the same span. Past a few dozen turns the -family alone exceeds the span attribute limit and evicts the core telemetry -written before it. The ceiling covers both directions together, since a budget -handed to each direction separately doubles. An eighth is the largest share -that still fits beside the tool ceiling and the core of every vocabulary at -once, request parameters, cost breakdown and identity included. The complete -conversation still rides the JSON blob attributes; only the per-index -convenience keys are capped. +An eighth is the largest share that still fits beside the tool ceiling and the core +of every vocabulary at once. The complete conversation still rides the JSON blobs. """ diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index 6491ab0f79f..16fbb242ebd 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -485,13 +485,7 @@ def _indexed_message_count(attributes, prefix): @pytest.mark.parametrize("turns", [60, 200]) def test_long_conversation_does_not_evict_core_attributes(turns): - """Per-message OpenInference attributes must never crowd core telemetry off the span. - - With content capture on, the OpenInference vocabulary spells every prompt and - response message out as two per-index attributes. A few dozen turns overruns - the OTel SDK's 128-attribute span limit, which evicts oldest-first, so the - ``gen_ai.*`` set written before it is what disappears. - """ + """Per-message OpenInference attributes must never crowd core telemetry off the span.""" span = _conversation_span(["genai", "openinference"], _conversation_payload(turns)) a = span.attributes @@ -522,12 +516,7 @@ def test_short_conversation_keeps_every_message_indexed(): def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit(monkeypatch): - """The per-index keys are the only untruncated copy once the SDK clips string values. - - Operators bound attribute sizes with ``OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT``, - which cuts the ``input.value`` blob short. The system prompt and the live turn - then have to survive as their own short keys, whatever the conversation length. - """ + """The system prompt and the live turn keep their own keys once the SDK clips ``input.value``.""" monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") payload = _conversation_payload(60) payload["messages"][0] = {"role": "system", "content": "be terse"} @@ -547,11 +536,7 @@ def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit def test_message_cap_is_shared_across_input_and_output(): - """One span-wide allowance covers both directions, and the response always keeps a share. - - A long prompt takes what a single reply leaves over, and a many-choice reply - cannot take the whole allowance away from the prompt either. - """ + """One span-wide allowance covers both directions, and the response always keeps a share.""" long_prompt = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=1)).attributes many_choices = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=20)).attributes @@ -569,13 +554,7 @@ def test_message_cap_is_shared_across_input_and_output(): def test_fully_populated_span_with_every_vocabulary_stays_within_the_attribute_limit(): - """Every capped family maxed at once still leaves the whole core intact. - - Every vocabulary in the registry plus ``legacy``, every request parameter, - every cost component, a hundred-plus tools, a two-hundred-turn prompt and - twenty choices is the worst case the two span-wide ceilings have to absorb - together. - """ + """Every capped family maxed at once still leaves the whole core intact.""" payload = _conversation_payload( 200, choices=20, From 107b4ec4db64985de0b3651f401b290ea09e81ed Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:12:39 +0000 Subject: [PATCH 036/297] fix(redis): log a timeout streak once per interval instead of one line per cache call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 140 ++++++++++++------ litellm/constants.py | 3 + tests/test_litellm/caching/test_dual_cache.py | 57 +++++++ .../test_litellm/caching/test_redis_cache.py | 57 +++++++ 4 files changed, 211 insertions(+), 46 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2c36995c4f8..eaac7ef7b0b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -15,6 +15,7 @@ import hashlib import inspect import json import logging +import threading import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar @@ -32,6 +33,7 @@ from litellm.constants import ( REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, + REDIS_TIMEOUT_LOG_INTERVAL, ) from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker @@ -404,13 +406,58 @@ class RedisCircuitBreakerOpenError(Exception): pass +class _RedisTimeoutLogThrottle: + """Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between.""" + + def __init__(self, interval: float, clock: Callable[[], float] = time.time) -> None: + self.interval = interval + self._clock = clock + self._lock = threading.Lock() + self._last_logged_at: float | None = None + self._suppressed = 0 + + def admit(self) -> int | None: + """Return the number of timeouts suppressed since the last admitted line, or None to suppress this one.""" + with self._lock: + now: Final = self._clock() + if self._last_logged_at is not None and now - self._last_logged_at < self.interval: + self._suppressed += 1 + return None + suppressed: Final = self._suppressed + self._suppressed = 0 + self._last_logged_at = now + return suppressed + + +_redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL) + + def log_redis_failure( logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False ) -> None: if isinstance(exc, RedisCircuitBreakerOpenError): - logger.debug("%s: %s", message, exc) + logger.debug("%s: %s", message, exc, stacklevel=2) return - logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None) + exc_info: Final = exc if with_traceback else None + if not _is_redis_timeout_failure(exc): + logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) + return + suppressed: Final = _redis_timeout_log_throttle.admit() + if suppressed is None: + logger.debug("%s: %s", message, exc, stacklevel=2) + return + if suppressed == 0: + logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) + return + logger.log( + level, + "%s: %s (%d more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + message, + exc, + suppressed, + exc_info=exc_info, + stacklevel=2, + ) @dataclass(frozen=True, slots=True) @@ -783,10 +830,8 @@ class RedisCache(BaseCache): ## LOGGING ## end_time = time.time() _duration = end_time - start_time - verbose_logger.error( - "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS", e ) raise e @@ -992,11 +1037,8 @@ class RedisCache(BaseCache): call_type=f"async_set_cache <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r", - str(e), - key, - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) raise e @@ -1044,10 +1086,8 @@ class RedisCache(BaseCache): event_metadata={"key": key}, ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) _record_swallowed_redis_failure(self._circuit_breaker, e) @@ -1094,7 +1134,6 @@ class RedisCache(BaseCache): start_time: Final = time.time() print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}") - cache_value: Final = None try: async with _redis_client.pipeline(transaction=False) as pipe: results: Final = await self._pipeline_helper(pipe, cache_list, ttl) @@ -1131,10 +1170,11 @@ class RedisCache(BaseCache): ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS %s, Writing value=%s", - str(e), - cache_value, + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", + e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) @@ -1177,10 +1217,8 @@ class RedisCache(BaseCache): ) ) # NON blocking - notify users Redis is throwing an exception - verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) raise e @@ -1216,10 +1254,11 @@ class RedisCache(BaseCache): ) ) # NON blocking - notify users Redis is throwing an exception - verbose_logger.error( - "LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS", + e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) @@ -1288,10 +1327,11 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async async_increment() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async async_increment() - Got exception from REDIS", + e, ) raise e @@ -1377,7 +1417,9 @@ class RedisCache(BaseCache): print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}") return self._get_cache_logic(cached_response=cached_response) except Exception as e: - verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e) + log_redis_failure( + verbose_logger, logging.ERROR, "litellm.caching.caching: get() - Got exception from REDIS", e + ) _record_swallowed_redis_failure(self._circuit_breaker, e) def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: @@ -1455,7 +1497,7 @@ class RedisCache(BaseCache): end_time=failed_at, parent_otel_span=parent_otel_span, ) - verbose_logger.error("Error occurred in batch get cache - %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in batch get cache", e) _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @@ -1574,7 +1616,7 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - verbose_logger.error("Error occurred in async batch get cache - %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in async batch get cache", e) _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @@ -1799,9 +1841,11 @@ class RedisCache(BaseCache): parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS %s", - str(e), + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS", + e, ) raise e @@ -1878,7 +1922,7 @@ class RedisCache(BaseCache): call_type=f"async_rpush <- {_get_call_stack_info()}", ) ) - verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e) raise e async def _pipeline_rpush_helper( @@ -1946,9 +1990,11 @@ class RedisCache(BaseCache): call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s", - str(e), + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS", + e, ) raise e @@ -2024,7 +2070,7 @@ class RedisCache(BaseCache): call_type=f"async_lpop <- {_get_call_stack_info()}", ) ) - verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache LPOP: - Got exception from REDIS", e) raise e async def _pipeline_lpop_helper( @@ -2135,8 +2181,10 @@ class RedisCache(BaseCache): call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s", - str(e), + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS", + e, ) raise e diff --git a/litellm/constants.py b/litellm/constants.py index 6b984c2673c..a32551b4480 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -459,6 +459,9 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED" # minimum seconds a timeout-only failure streak must span before it can open the breaker, # so one event-loop stall timing out many queued calls at once does not trip it REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)) +# seconds between Redis timeout log lines: the first timeout of a streak logs at the caller's level, +# later ones log at DEBUG until the interval passes and one line summarizes how many were suppressed +REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0")) # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart # (e.g. ElastiCache Serverless maintenance) is not reused while broken diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 4c9068722b8..850fa14106b 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -704,3 +704,60 @@ async def test_open_breaker_keeps_async_batch_read_memory_hits_and_releases_rese assert list(await cache.async_batch_get_cache(["k1", "k2"])) == ["v1", None] assert "k2" not in cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplog, monkeypatch): + """The in-memory fallback WARNING must not repeat for every timed-out increment during a blip. + + The rate limiter's pipeline increments and the dual cache increments each logged a WARNING per + call while Redis timed out, hundreds of lines per second before the breaker opened. The first + timeout of a streak keeps its WARNING, the rest are DEBUG until the summary interval passes. + """ + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + + clock = MagicMock(return_value=1_000.0) + monkeypatch.setattr( + redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + ) + + class _TimingOutRedis: + async def async_increment_pipeline(self, increment_list, **kwargs): + raise RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + + async def async_increment(self, key, value, **kwargs): + raise RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + + cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_TimingOutRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + increments = [RedisPipelineIncrementOperation(key="k", increment_value=1.0, ttl=60)] + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + for _ in range(100): + await cache.async_increment_cache_pipeline(increment_list=increments) + await cache.async_increment_cache("k", 1.0) + + visible = [r for r in caplog.records if r.levelno >= logging.WARNING] + assert [(r.levelno, r.getMessage()) for r in visible] == [ + ( + logging.WARNING, + "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" + " Timeout reading from 127.0.0.1:6379", + ) + ] + assert visible[0].filename == "dual_cache.py" + assert sum("Timeout reading from" in r.getMessage() for r in caplog.records) == 200 + + caplog.clear() + clock.return_value += 5.0 + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await cache.async_increment_cache("k", 1.0) + assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ + ( + logging.WARNING, + "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" + " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + ) + ] diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index bcae33b976e..d0974b2420c 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1202,3 +1202,60 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new new_probe_release.set() assert await new_probe == "new probe" assert breaker._state == breaker.CLOSED + + +def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_batch_redis_cache, caplog, monkeypatch): + """A Redis latency blip must not write one ERROR line per timed-out cache call. + + Before the breaker opens (up to REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION of timeouts) every + cache operation logged its own ERROR or WARNING line, so one single-worker proxy wrote + ~1100 lines in 5 s at LITELLM_LOG=WARNING. A timeout streak now logs its first failure, then + one summary line per REDIS_TIMEOUT_LOG_INTERVAL carrying the count of suppressed timeouts, + while every timeout stays visible at DEBUG. Hard connectivity failures keep their per-call line. + """ + import logging + + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + + clock = MagicMock(return_value=1_000.0) + monkeypatch.setattr( + redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + ) + sync_batch_redis_cache.redis_client.get.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + sync_batch_redis_cache.redis_client.mget.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + for _ in range(200): + assert sync_batch_redis_cache.get_cache("lit7520") is None + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7520"]) == {} + + timeout_records = [r for r in caplog.records if "Timeout reading from" in r.getMessage()] + assert len(timeout_records) == 201, "every timeout must stay visible at DEBUG" + assert [r.getMessage() for r in timeout_records if r.levelno >= logging.WARNING] == [ + "litellm.caching.caching: get() - Got exception from REDIS: Timeout reading from 127.0.0.1:6379" + ] + assert timeout_records[0].levelno == logging.ERROR + assert timeout_records[0].filename == "redis_cache.py" + assert timeout_records[0].lineno != timeout_records[-1].lineno, "the record must point at the cache operation" + + caplog.clear() + clock.return_value += 5.0 + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7520"]) == {} + assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ + ( + logging.ERROR, + "Error occurred in batch get cache: Timeout reading from 127.0.0.1:6379" + " (200 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + ) + ] + + caplog.clear() + sync_batch_redis_cache.redis_client.get.side_effect = OSError("redis unavailable") + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + for _ in range(3): + assert sync_batch_redis_cache.get_cache("lit7520") is None + assert [r.levelno for r in caplog.records if "redis unavailable" in r.getMessage()] == [logging.ERROR] * 3 From 9c84e98fb22bd0f6e2c359f335bbc329181bb8bd Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:25:32 +0000 Subject: [PATCH 037/297] fix(proxy): treat a Redis timeout in spend counter increments as an already-logged cache failure The cost tracking callback logged its own ERROR with a traceback for every request whose spend counter increment timed out, on top of the cache layer's throttled line. Timeouts now take the same path as breaker-open refusals: invalidate the counters and return. Also exposes is_redis_timeout_failure publicly for that caller and drops the comment on the new constant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 10 +++---- litellm/constants.py | 2 -- litellm/proxy/proxy_server.py | 4 +-- .../test_litellm/caching/test_redis_cache.py | 26 +++++++++---------- .../proxy/proxy_server/test_spend_counters.py | 22 ++++++++++++++++ 5 files changed, 42 insertions(+), 22 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index eaac7ef7b0b..7a3e689a667 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -330,7 +330,7 @@ def _redis_timeout_error_types() -> tuple[type, ...]: return (RedisTimeoutError, TimeoutError) -def _is_redis_timeout_failure(exc: BaseException) -> bool: +def is_redis_timeout_failure(exc: BaseException) -> bool: return isinstance(exc, _redis_timeout_error_types()) @@ -398,7 +398,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep """ if not _is_redis_health_failure(exc): return - breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc)) + breaker.record_failure(is_timeout=is_redis_timeout_failure(exc)) _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) @@ -439,7 +439,7 @@ def log_redis_failure( logger.debug("%s: %s", message, exc, stacklevel=2) return exc_info: Final = exc if with_traceback else None - if not _is_redis_timeout_failure(exc): + if not is_redis_timeout_failure(exc): logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) return suppressed: Final = _redis_timeout_log_throttle.admit() @@ -504,7 +504,7 @@ async def _run_under_circuit_breaker( result: Final = await call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) + breaker.record_failure(is_timeout=is_redis_timeout_failure(e)) raise _exit_circuit_breaker(breaker, admission) return result @@ -521,7 +521,7 @@ def _run_under_circuit_breaker_sync( result: Final = call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) + breaker.record_failure(is_timeout=is_redis_timeout_failure(e)) raise _exit_circuit_breaker(breaker, admission) return result diff --git a/litellm/constants.py b/litellm/constants.py index a32551b4480..60e1c682238 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -459,8 +459,6 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED" # minimum seconds a timeout-only failure streak must span before it can open the breaker, # so one event-loop stall timing out many queued calls at once does not trip it REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)) -# seconds between Redis timeout log lines: the first timeout of a streak logs at the caller's level, -# later ones log at DEBUG until the interval passes and one line summarizes how many were suppressed REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0")) # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c8fd7edfed6..318ea96dbcb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -251,7 +251,7 @@ import litellm._redis from litellm import Router from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, @@ -3411,7 +3411,7 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) except Exception as e: await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) - if isinstance(e, RedisCircuitBreakerOpenError): + if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e): return raise for item, current_value in zip(pending, results or ()): diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index d0974b2420c..5840f450ac6 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -977,17 +977,17 @@ async def test_stale_timeout_does_not_let_sub_threshold_hard_failures_open_the_b from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) await asyncio.sleep(0.06) for _ in range(breaker.failure_threshold - 1): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) assert breaker.is_open() is False, "2 hard failures and 1 stale timeout are below both thresholds" - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) assert breaker.is_open() is True, "the threshold-th hard failure must still open it" @@ -999,19 +999,19 @@ async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) await asyncio.sleep(0.06) for _ in range(breaker.failure_threshold): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) assert breaker.is_open() is False, "the burst is instantaneous, so the duration gate must hold it closed" await asyncio.sleep(0.06) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) assert breaker.is_open() is True, "the same run of timeouts persisting past the duration must open it" @@ -1022,7 +1022,7 @@ async def test_breaker_metrics_track_state_and_failure_class(): from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure def sample(name, labels=None): return REGISTRY.get_sample_value(name, labels) or 0.0 @@ -1034,9 +1034,9 @@ async def test_breaker_metrics_track_state_and_failure_class(): closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("t"))) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("t"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) == timeout_before + 1 assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) == hard_before + 2 diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 2f47736a398..19b5a11af33 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1180,6 +1180,28 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur fake_cache.in_memory_cache.set_cache.assert_not_called() +@pytest.mark.asyncio +async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_returns(monkeypatch): + """A Redis timeout is the streak the breaker is already counting and the cache layer already logged. + + Re-raising it sent every request in the pre-open window through the cost callback's error + path, which logged a traceback and fired the failed-tracking alert once per request. + """ + from redis.exceptions import TimeoutError as RedisTimeoutError + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_increment_pipeline = AsyncMock( + side_effect=RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + ) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + await ps._apply_spend_counter_increments(_two_pending_increments()) + + deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) + assert deleted_keys == ["spend:key:k", "spend:team:t"] + fake_cache.in_memory_cache.set_cache.assert_not_called() + + @pytest.mark.asyncio async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_raises(monkeypatch): fake_cache = _make_spend_counter_cache() From f681a978f06baa13da0f0c24f7b1ac3a20d9a02a Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:37:36 +0000 Subject: [PATCH 038/297] fix(redis): use a monotonic clock for the timeout log throttle and trim test docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 2 +- tests/test_litellm/caching/test_dual_cache.py | 7 +------ tests/test_litellm/caching/test_redis_cache.py | 9 +-------- .../proxy/proxy_server/test_spend_counters.py | 6 +----- 4 files changed, 4 insertions(+), 20 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 7a3e689a667..e5e27d1e02b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -409,7 +409,7 @@ class RedisCircuitBreakerOpenError(Exception): class _RedisTimeoutLogThrottle: """Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between.""" - def __init__(self, interval: float, clock: Callable[[], float] = time.time) -> None: + def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None: self.interval = interval self._clock = clock self._lock = threading.Lock() diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 850fa14106b..6f29be00b30 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -708,12 +708,7 @@ async def test_open_breaker_keeps_async_batch_read_memory_hits_and_releases_rese @pytest.mark.asyncio async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplog, monkeypatch): - """The in-memory fallback WARNING must not repeat for every timed-out increment during a blip. - - The rate limiter's pipeline increments and the dual cache increments each logged a WARNING per - call while Redis timed out, hundreds of lines per second before the breaker opened. The first - timeout of a streak keeps its WARNING, the rest are DEBUG until the summary interval passes. - """ + """The first fallback WARNING of a timeout streak logs, the rest stay at DEBUG until the summary.""" from redis.exceptions import TimeoutError as RedisTimeoutError from litellm.caching import redis_cache as redis_cache_module diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 5840f450ac6..4ca33894aed 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1205,14 +1205,7 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_batch_redis_cache, caplog, monkeypatch): - """A Redis latency blip must not write one ERROR line per timed-out cache call. - - Before the breaker opens (up to REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION of timeouts) every - cache operation logged its own ERROR or WARNING line, so one single-worker proxy wrote - ~1100 lines in 5 s at LITELLM_LOG=WARNING. A timeout streak now logs its first failure, then - one summary line per REDIS_TIMEOUT_LOG_INTERVAL carrying the count of suppressed timeouts, - while every timeout stays visible at DEBUG. Hard connectivity failures keep their per-call line. - """ + """A timeout streak logs its first failure plus one summary per interval; other failures log per call.""" import logging from redis.exceptions import TimeoutError as RedisTimeoutError diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 19b5a11af33..4a1fc389d3e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1182,11 +1182,7 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur @pytest.mark.asyncio async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_returns(monkeypatch): - """A Redis timeout is the streak the breaker is already counting and the cache layer already logged. - - Re-raising it sent every request in the pre-open window through the cost callback's error - path, which logged a traceback and fired the failed-tracking alert once per request. - """ + """A Redis timeout invalidates the counters and returns without reaching the cost callback's error path.""" from redis.exceptions import TimeoutError as RedisTimeoutError fake_cache = _make_spend_counter_cache() From 28f2d1f0168aa31639a23447d391516129267069 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:58:36 +0000 Subject: [PATCH 039/297] test(redis): cover the write and list timeout paths going through the shared log throttle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_litellm/caching/test_redis_cache.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 4ca33894aed..5ea21dae539 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1252,3 +1252,60 @@ def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_bat for _ in range(3): assert sync_batch_redis_cache.get_cache("lit7520") is None assert [r.levelno for r in caplog.records if "redis unavailable" in r.getMessage()] == [logging.ERROR] * 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_method", + [ + pytest.param(lambda c: c.async_set_cache_pipeline([("lit7520", "v")]), id="async_set_cache_pipeline"), + pytest.param(lambda c: c.async_set_cache_sadd("lit7520", ["v"], ttl=None), id="async_set_cache_sadd"), + pytest.param(lambda c: c.async_increment("lit7520", 1.0), id="async_increment"), + pytest.param( + lambda c: c.async_increment_pipeline([{"key": "lit7520", "increment_value": 1.0, "ttl": 60}]), + id="async_increment_pipeline", + ), + pytest.param(lambda c: c.async_rpush("lit7520", ["v"]), id="async_rpush"), + pytest.param( + lambda c: c.async_rpush_pipeline([{"key": "lit7520", "values": ["v"]}]), id="async_rpush_pipeline" + ), + pytest.param(lambda c: c.async_lpop("lit7520"), id="async_lpop"), + pytest.param(lambda c: c.async_lpop_pipeline([{"key": "lit7520", "count": 1}]), id="async_lpop_pipeline"), + ], +) +async def test_write_path_timeouts_inside_the_interval_stay_at_debug(call_method, caplog, monkeypatch, redis_no_ping): + """A write or list operation timing out mid-streak is counted by the throttle instead of logging its own ERROR.""" + import contextlib + import logging + + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + + clock = MagicMock(return_value=1_000.0) + throttle = _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + assert throttle.admit() == 0 + monkeypatch.setattr(redis_cache_module, "_redis_timeout_log_throttle", throttle) + + timeout = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + client = MagicMock() + client.pipeline.return_value.__aenter__.side_effect = timeout + client.sadd = AsyncMock(side_effect=timeout) + client.incrbyfloat = AsyncMock(side_effect=timeout) + client.rpush = AsyncMock(side_effect=timeout) + client.lpop = AsyncMock(side_effect=timeout) + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + cache = RedisCache() + + with ( + patch.object(cache, "init_async_client", return_value=client), + caplog.at_level(logging.DEBUG, logger="LiteLLM"), + ): + with contextlib.suppress(RedisTimeoutError): + await call_method(cache) + + timeout_records = [r for r in caplog.records if "Timeout reading from" in r.getMessage()] + assert [(r.levelno, r.filename) for r in timeout_records] == [(logging.DEBUG, "redis_cache.py")] + clock.return_value += 5.0 + assert throttle.admit() == 1 From 88de192dcf55e26f0f2cabb4d769a825dd8dbc7f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:29:04 -0700 Subject: [PATCH 040/297] test: bind management E2E callers and isolate JWT actors --- .github/e2e-stack/assert_tests_ran.py | 15 + .github/e2e-stack/oidc-profile.sh | 5 + .github/e2e-stack/select_tests.py | 2 + .github/workflows/test-e2e-changed.yml | 2 +- .../test_e2e_changed_gate.py | 23 ++ tests/e2e/CONTRIBUTING.md | 6 +- .../e2e/coverage_registry/management_cases.py | 151 +++++++++ tests/e2e/coverage_registry/mgmt.yaml | 8 + tests/e2e/e2e_http.py | 31 +- tests/e2e/idp.py | 254 ++++++++++++++- tests/e2e/junit_properties.py | 3 +- tests/e2e/management/conftest.py | 19 +- tests/e2e/management/jwt_actors.py | 175 ++++++++++ tests/e2e/management/management_client.py | 122 ++++--- .../e2e/management/test_jwt_management_e2e.py | 233 +++++++++++-- tests/e2e/models.py | 28 +- tests/e2e/proxy_client.py | 116 ++++--- tests/e2e/test_e2e_http.py | 10 + tests/e2e/test_idp.py | 122 ++++++- tests/e2e/test_proxy_client.py | 307 +++++++++++++++++- tests/e2e/transport.py | 17 +- tests/e2e/ui/oidcSetup.ts | 30 ++ tests/e2e/ui/playwright.oidc.config.ts | 22 ++ 23 files changed, 1541 insertions(+), 160 deletions(-) create mode 100755 .github/e2e-stack/oidc-profile.sh create mode 100644 tests/e2e/coverage_registry/management_cases.py create mode 100644 tests/e2e/management/jwt_actors.py create mode 100644 tests/e2e/ui/oidcSetup.ts create mode 100644 tests/e2e/ui/playwright.oidc.config.ts diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index c4348c20873..2303c42f4fb 100644 --- a/.github/e2e-stack/assert_tests_ran.py +++ b/.github/e2e-stack/assert_tests_ran.py @@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET from pathlib import Path from typing import Final +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e")) +from coverage_registry.management_cases import MANAGEMENT_CASES + def main() -> int: selected: Final = tuple(sys.argv[2:]) @@ -16,6 +19,17 @@ def main() -> int: case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) ) missing: Final = tuple(path for path in selected if path not in passed) + required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected) + passed_nodes: Final = frozenset( + prop.get("value") + for case in cases + if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) + for prop in case.findall("./properties/property") + if prop.get("name") == "management_node" + ) + missing_nodes: Final = required_nodes - passed_nodes + for node in sorted(missing_nodes): + _ = sys.stdout.write(f"::error::required management case did not pass: {node}\n") for path in selected: collected: Final = sum(case.get("file") == path for case in cases) skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) @@ -27,6 +41,7 @@ def main() -> int: if ( selected and not missing + and not missing_nodes and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error")) ): return 0 diff --git a/.github/e2e-stack/oidc-profile.sh b/.github/e2e-stack/oidc-profile.sh new file mode 100755 index 00000000000..84eaaaf8051 --- /dev/null +++ b/.github/e2e-stack/oidc-profile.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +set -euo pipefail +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "${REPO_ROOT}" +exec uv run --no-sync python tests/e2e/idp.py "$@" diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 238818a0d36..982e93cf642 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile( HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" r"|^tests/e2e/idp_realm\.json$" + r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$" + r"|^tests/e2e/coverage_registry/management_cases\.py$" r"|^tests/e2e/gateway/" r"|^\.github/e2e-stack/" r"|^\.github/workflows/test-e2e-changed\.yml$" diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 1db597ff673..c9f08deb36e 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -183,7 +183,7 @@ jobs: log="${RUNNER_TEMP}/e2e-pass-${pass}.log" echo "::group::pass ${pass} of 3" set +e - uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \ + uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \ -o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1 status=$? uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}" diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 588402e3996..101816c7f11 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -81,6 +81,25 @@ def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None assert result.returncode == 1 +@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer")) +def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None: + suite: Final = ET.Element("testsuite") + path: Final = "tests/e2e/management/test_jwt_management_e2e.py" + case: Final = ET.SubElement(suite, "testcase", file=path) + properties: Final = ET.SubElement(case, "properties") + _ = ET.SubElement( + properties, + "property", + name="management_node", + value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]", + ) + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True) + assert result.returncode == 1 + assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout + + def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_path: Path) -> None: env_path: Final = tmp_path / ".env" @@ -141,6 +160,10 @@ def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them( ( "tests/e2e/proxy_client.py", "tests/e2e/conftest.py", + "tests/e2e/management/management_client.py", + "tests/e2e/management/jwt_actors.py", + "tests/e2e/management/conftest.py", + "tests/e2e/coverage_registry/management_cases.py", "tests/e2e/pytest.ini", "tests/e2e/gateway/stage_mirror_ci_config.yml", ".github/e2e-stack/up.sh", diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 44564a51e26..78c05ea4b30 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -60,7 +60,11 @@ The suites run against a live proxy, so bring one up first by running the litell Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack. - Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write. + Management tests can bind a credential once with `client.with_caller(Caller(...))`; direct calls, delegated helpers and replica read-backs then retain that caller. Explicit `caller_key` arguments override the binding. Keep the original master-backed client for bootstrap and cleanup. `actor_factory` lazily provisions database roles and tenant memberships, with `database_role` tokens carrying no groups and `group_scoped` actors retaining the existing team route gate. Token minting is explicit through `actor.mint_caller(idp)`. The factory runs requests without backend retries and reports cleanup failures. `coverage_registry/management_cases.py` records exact canary nodes and non-secret actor labels; the CI execution assertion rejects a missing or skipped actor row + + For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" `. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts + + `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`: diff --git a/tests/e2e/coverage_registry/management_cases.py b/tests/e2e/coverage_registry/management_cases.py new file mode 100644 index 00000000000..15dc7d333c5 --- /dev/null +++ b/tests/e2e/coverage_registry/management_cases.py @@ -0,0 +1,151 @@ +from dataclasses import dataclass +from typing import Final, Literal + +CredentialKind = Literal["master", "idp_admin", "direct_jwt", "virtual_key", "dashboard_session"] +DependencyProfile = Literal["management_only", "real_oidc_browser", "external_provider_required"] + + +@dataclass(frozen=True, slots=True) +class ManagementCase: + node: str + credential_kind: CredentialKind + actor: str + profile: str + method: Literal["GET", "POST"] + path: str + operation_family: str + dependency_profile: DependencyProfile = "management_only" + + +JWT_FILE: Final = "tests/e2e/management/test_jwt_management_e2e.py" +JWT_CLASS: Final = f"{JWT_FILE}::TestJwtManagement" +ACTORS: Final = ( + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", +) +MANAGEMENT_CASES: Final = tuple( + ManagementCase( + node=f"{JWT_CLASS}::test_actor_subject_and_database_role[{role}]", + credential_kind="direct_jwt", + actor=role, + profile="database_role", + method="GET", + path="/user/info", + operation_family="identity", + ) + for role in ACTORS +) + ( + ManagementCase( + node=f"{JWT_CLASS}::test_admin_viewer_reads_but_cannot_update", + credential_kind="direct_jwt", + actor="proxy_admin_viewer", + profile="database_role", + method="POST", + path="/key/update", + operation_family="denial", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[direct_jwt]", + credential_kind="direct_jwt", + actor="proxy_admin", + profile="group_scoped", + method="POST", + path="/key/generate", + operation_family="key_lifecycle", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]", + credential_kind="virtual_key", + actor="proxy_admin", + profile="database_role", + method="POST", + path="/key/generate", + operation_family="key_lifecycle", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_two_actor_sets_keep_tenants_and_keys_isolated", + credential_kind="direct_jwt", + actor="team_member", + profile="group_scoped", + method="GET", + path="/key/info", + operation_family="tenant_isolation", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_member_cannot_write_and_another_team_cannot_read_the_key", + credential_kind="direct_jwt", + actor="team_member", + profile="group_scoped", + method="POST", + path="/key/update", + operation_family="tenant_isolation", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_multi_group_actor_keeps_exact_memberships", + credential_kind="master", + actor="bootstrap", + profile="group_scoped", + method="GET", + path="/team/info", + operation_family="memberships", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_successful_actor_cleanup_removes_owned_state", + credential_kind="master", + actor="bootstrap", + profile="failure_cleanup", + method="GET", + path="/team/info", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[group]", + credential_kind="idp_admin", + actor="idp_admin", + profile="failure_cleanup", + method="POST", + path="/groups", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[user]", + credential_kind="idp_admin", + actor="idp_admin", + profile="failure_cleanup", + method="POST", + path="/users", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_oidc_browser_profile_identity_mapping", + credential_kind="direct_jwt", + actor="internal_user", + profile="oidc_configuration", + method="GET", + path="/protocol/openid-connect/userinfo", + operation_family="oidc_identity", + ), +) + + +def canonical_node(node: str) -> str: + return node if node.startswith("tests/e2e/") else f"tests/e2e/{node}" + + +def case_properties(node: str) -> tuple[tuple[str, str], ...]: + case: Final = next((case for case in MANAGEMENT_CASES if case.node == canonical_node(node)), None) + if case is None: + return () + return ( + ("management_node", case.node), + ("credential_kind", case.credential_kind), + ("actor", case.actor), + ("auth_profile", case.profile), + ("dependency_profile", case.dependency_profile), + ) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d1227fe7c0c..31ad61ba3e2 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -90,3 +90,11 @@ - {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} - {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} - {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} + +- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"} +- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"} +- {id: mgmt.user.oidc.identity_mapping, module: mgmt, tier: P0, surface: api, assertions: [identity_mapping], source: "tests/e2e/idp.py", rationale: "IdP configuration canary only: confidential-client token and userinfo subjects match the seeded user; application SSO is separate"} +- {id: mgmt.team.jwt.tenant_isolation, module: mgmt, tier: P0, surface: api, assertions: [tenant_isolation], source: "auth/handle_jwt.py", rationale: "Isolated team actors read their own key and receive 403 for the other tenant key"} +- {id: mgmt.team.jwt.multiple_memberships, module: mgmt, tier: P0, surface: api, assertions: [multiple_memberships], source: "auth/handle_jwt.py", rationale: "A multi-group actor has exactly the configured memberships without admin scope"} +- {id: mgmt.user.jwt.cleanup, module: mgmt, tier: P0, surface: api, assertions: [cleanup], source: "management_endpoints/internal_user_endpoints.py", rationale: "Owned users teams organizations keys and IdP objects disappear after successful cleanup"} +- {id: mgmt.user.jwt.partial_cleanup, module: mgmt, tier: P0, surface: api, assertions: [partial_cleanup], source: "auth/handle_jwt.py", rationale: "Partial identity setup removes the group and user created before failure"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index ce069720c6e..67370c98274 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -16,9 +16,11 @@ requests itself imports. from __future__ import annotations import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Generator, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass -from typing import Final, Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast +from typing import Final, Generic, Literal, NewType, Protocol, TypeVar, cast import pytest import requests @@ -36,8 +38,8 @@ class Headers(BaseModel): class AuthHeaders(Headers): # litellm accepts either; set whichever the call needs, leave the other None. - authorization: str | None = None - x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key") + authorization: str | None = Field(default=None, repr=False) + x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key", repr=False) class AnthropicHeaders(AuthHeaders): @@ -292,6 +294,22 @@ def _params(params: BaseModel | None) -> dict[str, str]: TRANSIENT_STATUSES: frozenset[int] = frozenset({529}) RETRY_ATTEMPTS: int = 3 +_QUALIFICATION: Final[ContextVar[bool]] = ContextVar("e2e_qualification", default=False) + + +def retry_attempts(default: int) -> int: + return 1 if _QUALIFICATION.get() else default + + +@contextmanager +def without_retries() -> Generator[None]: + token: Final = _QUALIFICATION.set(True) + try: + yield + finally: + _QUALIFICATION.reset(token) + + RETRY_BACKOFF_SECONDS: float = 0.5 @@ -319,7 +337,7 @@ def request_with_retry[T: RetryableResponse]( hang should surface as a hang instead of doubling the wall clock. Every retry prints, so flakiness stays visible in the run log instead of vanishing into green.""" - for attempt in range(1, RETRY_ATTEMPTS): + for attempt in range(1, retry_attempts(RETRY_ATTEMPTS)): resp = issue() if resp.status_code not in TRANSIENT_STATUSES: return resp @@ -414,6 +432,7 @@ def get_external[R: BaseModel]( url: str, *, response_type: type[R], + headers: BaseModel | None = None, timeout: float = 30.0, ) -> Result[R]: """GET an absolute URL outside the proxy (e.g. a public /.well-known document). @@ -422,7 +441,7 @@ def get_external[R: BaseModel]( try: resp = requests.get( url, - headers={"Accept": "application/json"}, + headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})}, timeout=timeout, ) except requests.RequestException as exc: diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 6d2fc84eb27..12db91bbd88 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -2,11 +2,17 @@ from __future__ import annotations +import base64 import os import secrets +import signal +import subprocess +import sys import warnings from collections.abc import Callable -from dataclasses import dataclass, field +from contextlib import ExitStack +from dataclasses import dataclass, field, replace +from types import FrameType from typing import Final, Literal import pytest @@ -14,11 +20,15 @@ from e2e_http import ( AuthHeaders, ExternalWrite, NetworkError, + NoBody, Result, Success, + UnknownApiError, delete_external, + get_external, post_form_external, post_json_external, + unwrap, ) from pydantic import BaseModel, Field @@ -46,7 +56,9 @@ class TokenGrantForm(BaseModel): grant_type: Literal["password"] = "password" client_id: str username: str - password: str + password: str = Field(repr=False) + client_secret: str | None = Field(default=None, repr=False) + scope: str | None = None class TokenResponse(BaseModel): @@ -63,7 +75,7 @@ class GroupCreateBody(BaseModel): class PasswordCredential(BaseModel): type: Literal["password"] = "password" - value: str + value: str = Field(repr=False) temporary: bool = False @@ -101,8 +113,20 @@ class Identity: user_id: str username: str password: str = field(repr=False) - group: str - group_id: str + groups: tuple[str, ...] + group_ids: tuple[str, ...] + + @property + def group(self) -> str: + if len(self.groups) != 1: + raise ValueError("A single-group identity is required") + return self.groups[0] + + @property + def group_id(self) -> str: + if len(self.group_ids) != 1: + raise ValueError("A single-group identity is required") + return self.group_ids[0] @dataclass(frozen=True, slots=True) @@ -111,6 +135,10 @@ class Keycloak: realm: str admin_username: str admin_password: str = field(repr=False) + strict_cleanup: bool = False + + def with_strict_cleanup(self) -> Keycloak: + return replace(self, strict_cleanup=True) @property def issuer(self) -> str: @@ -150,7 +178,9 @@ class Keycloak: f"group {name}", ) - def create_user(self, *, username: str, email: str, password: str, group: str) -> str: + def create_user( + self, *, username: str, email: str, password: str, group: str | None = None, groups: tuple[str, ...] = () + ) -> str: return created_id( post_json_external( self._admin_url("/users"), @@ -158,7 +188,7 @@ class Keycloak: json=UserCreateBody( username=username, email=email, - groups=(group,), + groups=(group,) if group is not None else groups, credentials=(PasswordCredential(value=password),), ), ), @@ -171,14 +201,28 @@ class Keycloak: def delete_group(self, group_id: str) -> None: self._delete(f"/groups/{group_id}") + def assert_absent(self, kind: Literal["users", "groups", "clients"], resource_id: str) -> None: + result: Final = get_external( + self._admin_url(f"/{kind}/{resource_id}"), + headers=self._admin_headers(), + response_type=NoBody, + ) + assert isinstance(result, UnknownApiError) and result.status_code == 404, ( + f"Owned IdP {kind} still exists: {result}" + ) + def _delete(self, path: str) -> None: try: headers: Final = self._admin_headers() except pytest.fail.Exception as exc: + if self.strict_cleanup: + raise RuntimeError(f"Keycloak cleanup could not authenticate for {path}") from exc warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2) return result: Final = delete_external(self._admin_url(path), headers=headers) if result.status_code not in (204, 404): + if self.strict_cleanup: + raise RuntimeError(f"Keycloak cleanup failed for {path}: HTTP {result.status_code}") warnings.warn( f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}", RuntimeWarning, @@ -188,15 +232,34 @@ class Keycloak: def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity: """Create `group` and a user in it, credentialed with a password generated for this test alone, and hand back the identity a token can be minted for.""" - group_id: Final = self.create_group(group) - defer(lambda: self.delete_group(group_id)) + return self.provision_groups(marker=marker, groups=(group,), defer=defer) + + def provision_groups( + self, *, marker: str, groups: tuple[str, ...], defer: Callable[[Callable[[], object]], None] + ) -> Identity: + def provision_group(name: str) -> str: + created: Final = self.create_group(name) + defer(lambda: self.delete_group(created)) + return created + + group_ids: Final = tuple(provision_group(group) for group in groups) + return self.provision_user(marker=marker, groups=groups, group_ids=group_ids, defer=defer) + + def provision_user( + self, + *, + marker: str, + groups: tuple[str, ...], + group_ids: tuple[str, ...], + defer: Callable[[Callable[[], object]], None], + ) -> Identity: username: Final = f"e2e-jwt-user-{marker}" password: Final = secrets.token_urlsafe(24) user_id: Final = self.create_user( - username=username, email=f"{username}@example.com", password=password, group=group + username=username, email=f"{username}@example.com", password=password, groups=groups ) defer(lambda: self.delete_user(user_id)) - return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id) + return Identity(user_id=user_id, username=username, password=password, groups=groups, group_ids=group_ids) def access_token( self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None @@ -211,6 +274,65 @@ class Keycloak: ) return self._token(result, f"a token for {identity.username}") + def discovery(self) -> Discovery: + return unwrap(get_external(f"{self.issuer}/.well-known/openid-configuration", response_type=Discovery)) + + def browser_client(self, *, callback_url: str, defer: Callable[[Callable[[], object]], None]) -> BrowserClient: + client: Final = BrowserClient( + client_id=f"e2e-browser-{secrets.token_hex(8)}", + secret=secrets.token_urlsafe(32), + callback_url=callback_url, + ) + resource_id: Final = created_id( + post_json_external( + self._admin_url("/clients"), + headers=self._admin_headers(), + json=BrowserClientBody( + clientId=client.client_id, + secret=client.secret, + redirectUris=(callback_url,), + ), + ), + "browser client", + ) + defer(lambda: self._delete(f"/clients/{resource_id}")) + configured: Final = unwrap( + get_external( + self._admin_url(f"/clients/{resource_id}"), + headers=self._admin_headers(), + response_type=BrowserClientBody, + ) + ) + assert configured.redirect_uris == (callback_url,) + assert configured.standard_flow_enabled and not configured.public_client + assert configured.attributes.pkce == "S256" + return client + + def browser_token(self, identity: Identity, client: BrowserClient) -> str: + return self._token( + post_form_external( + self.token_url(self.realm), + form=TokenGrantForm( + client_id=client.client_id, + client_secret=client.secret, + username=identity.username, + password=identity.password, + scope="openid email", + ), + response_type=TokenResponse, + ), + "browser-profile identity mapping", + ) + + def userinfo(self, token: str) -> UserInfo: + return unwrap( + get_external( + f"{self.issuer}/protocol/openid-connect/userinfo", + headers=AuthHeaders(authorization=f"Bearer {token}"), + response_type=UserInfo, + ) + ) + def keycloak_from_env() -> Keycloak: admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip() @@ -226,3 +348,113 @@ def keycloak_from_env() -> Keycloak: admin_username=admin_username, admin_password=admin_password, ) + + +class TokenClaims(BaseModel): + sub: str + iss: str + aud: str | tuple[str, ...] + exp: int + scope: str = "" + groups: tuple[str, ...] = () + + +class Discovery(BaseModel): + issuer: str + authorization_endpoint: str + token_endpoint: str + userinfo_endpoint: str + jwks_uri: str + + +class UserInfo(BaseModel): + sub: str + email: str + + +class BrowserAttributes(BaseModel): + pkce: str = Field(default="S256", alias="pkce.code.challenge.method") + + +class AudienceConfig(BaseModel): + audience: str = Field(default="litellm-e2e", alias="included.custom.audience") + access_token: str = Field(default="true", alias="access.token.claim") + id_token: str = Field(default="false", alias="id.token.claim") + + +class AudienceMapper(BaseModel): + name: str = "litellm-audience" + protocol: str = "openid-connect" + mapper: str = Field(default="oidc-audience-mapper", alias="protocolMapper") + config: AudienceConfig = Field(default_factory=AudienceConfig) + + +class BrowserClientBody(BaseModel): + client_id: str = Field(alias="clientId") + secret: str = Field(repr=False) + redirect_uris: tuple[str, ...] = Field(alias="redirectUris") + enabled: bool = True + public_client: bool = Field(default=False, alias="publicClient") + standard_flow_enabled: bool = Field(default=True, alias="standardFlowEnabled") + direct_access_grants_enabled: bool = Field(default=True, alias="directAccessGrantsEnabled") + default_client_scopes: tuple[str, ...] = Field(default=("email", "basic"), alias="defaultClientScopes") + attributes: BrowserAttributes = Field(default_factory=BrowserAttributes) + protocol_mappers: tuple[AudienceMapper, ...] = Field(default=(AudienceMapper(),), alias="protocolMappers") + + +@dataclass(frozen=True, slots=True) +class BrowserClient: + client_id: str + secret: str = field(repr=False) + callback_url: str + + def environment(self, discovery: Discovery) -> dict[str, str]: + return { + "GENERIC_CLIENT_ID": self.client_id, + "GENERIC_CLIENT_SECRET": self.secret, + "GENERIC_USER_ID_ATTRIBUTE": "sub", + "GENERIC_AUTHORIZATION_ENDPOINT": discovery.authorization_endpoint, + "GENERIC_TOKEN_ENDPOINT": discovery.token_endpoint, + "GENERIC_USERINFO_ENDPOINT": discovery.userinfo_endpoint, + "GENERIC_CLIENT_USE_PKCE": "true", + "GENERIC_SCOPE": "openid email", + } + + +def token_claims(token: str) -> TokenClaims: + payload: Final = token.split(".")[1] + return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) + + +def run_oidc_profile(proxy_url: str, command: list[str]) -> int: + idp: Final = keycloak_from_env().with_strict_cleanup() + with ExitStack() as cleanup: + + def terminate(signum: int, frame: FrameType | None) -> None: + raise SystemExit(128 + signum) + + previous: Final = signal.signal(signal.SIGTERM, terminate) + cleanup.callback(signal.signal, signal.SIGTERM, previous) + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer) + environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url} + with subprocess.Popen(command, env=environment) as child: + try: + return child.wait() + finally: + if child.poll() is None: + child.terminate() + try: + child.wait(timeout=5) + except subprocess.TimeoutExpired: + child.kill() + child.wait() + + +if __name__ == "__main__": + if len(sys.argv) < 3: + raise SystemExit("Usage: idp.py PROXY_URL COMMAND [ARG ...]; requires a running test IdP") + raise SystemExit(run_oidc_profile(sys.argv[1], sys.argv[2:])) diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index c5971c5362c..b9f5da871ae 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -19,6 +19,7 @@ from __future__ import annotations from collections.abc import Iterable import pytest +from coverage_registry.management_cases import case_properties # Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing # at runtime names this suite's place in the repo. test_junit_properties.py @@ -94,7 +95,7 @@ def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: ("package", package_from_nodeid(item.nodeid)), ("covers", ",".join(covers_from_item(item))), ("source", source_from_item(item)), - ) + ) + case_properties(item.nodeid) def attach_result_properties(item: pytest.Item) -> None: diff --git a/tests/e2e/management/conftest.py b/tests/e2e/management/conftest.py index bd69c8c0ff3..5a11b634085 100644 --- a/tests/e2e/management/conftest.py +++ b/tests/e2e/management/conftest.py @@ -5,8 +5,14 @@ holds the shared ProxyClient so `resources` / `scoped_key` clean up keys, teams, users, and orgs this suite creates. """ -import pytest +from collections.abc import Generator +from typing import Final +import pytest +from e2e_http import without_retries +from idp import Keycloak +from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory from management_client import ManagementClient, build_client from proxy_client import ProxyClient @@ -21,3 +27,14 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> ManagementClient: return build_client(proxy) + + +@pytest.fixture +def actor_factory(proxy: ProxyClient, idp: Keycloak) -> Generator[ActorFactory]: + bootstrap: Final = build_client(proxy) + resources: Final = ResourceManager(client=proxy, strict_cleanup=True) + with without_retries(): + try: + yield ActorFactory(bootstrap=bootstrap, idp=idp, resources=resources) + finally: + resources.teardown() diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py new file mode 100644 index 00000000000..909d1652ada --- /dev/null +++ b/tests/e2e/management/jwt_actors.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final, Literal + +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from idp import ADMIN_CLIENT_ID, TESTS_CLIENT_ID, Identity, Keycloak +from lifecycle import ResourceManager +from management.management_client import ManagementClient +from models import ( + KeyGenerateBody, + KeyGenerateResponse, + OrgDeleteBody, + OrgDeleteResponse, + OrgMemberAddBody, + OrgMemberEntry, + OrgNewBody, + TeamDeleteBody, + TeamMemberAddBody, + TeamMemberEntry, + TeamNewBody, + UserNewBody, + UserRole, +) +from proxy_client import Caller + +ActorRole = Literal[ + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", +] +ActorProfile = Literal["database_role", "group_scoped"] + + +@dataclass(frozen=True, slots=True) +class Tenant: + organization_id: str + team_id: str + group_id: str + + +@dataclass(frozen=True, slots=True) +class Actor: + identity: Identity + role: ActorRole + global_role: UserRole + profile: ActorProfile + tenants: tuple[Tenant, ...] + + def mint_caller(self, idp: Keycloak) -> Caller: + return Caller( + credential=idp.access_token( + self.identity, client_id=ADMIN_CLIENT_ID if self.role == "proxy_admin" else TESTS_CLIENT_ID + ), + kind="direct_jwt", + role=self.role, + tenant=self.tenants[0].organization_id if self.tenants else None, + ) + + +@dataclass(frozen=True, slots=True) +class ActorFactory: + bootstrap: ManagementClient + idp: Keycloak + resources: ResourceManager + + def __post_init__(self) -> None: + if self.bootstrap.proxy.caller is not None: + raise ValueError("Actor bootstrap requires a separately held master client") + + def key(self, tenant: Tenant | None = None, *, user_id: str | None = None) -> KeyGenerateResponse: + created: Final = unwrap( + self.bootstrap.generate_key( + KeyGenerateBody( + team_id=tenant.team_id if tenant is not None else None, + user_id=user_id, + key_alias=f"e2e-actor-key-{unique_marker()}", + ) + ) + ) + self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key)) + return created + + def tenant(self) -> Tenant: + marker: Final = unique_marker() + organization_id: Final = self.bootstrap.create_org(OrgNewBody(organization_alias=f"e2e-organization-{marker}")) + self.resources.defer( + lambda: unwrap( + self.bootstrap.proxy.transport.delete( + "/organization/delete", + headers=self.bootstrap.proxy.management_headers(), + json=OrgDeleteBody(organization_ids=[organization_id]), + response_type=OrgDeleteResponse, + ) + ) + ) + team_id: Final = self.bootstrap.proxy.create_team( + TeamNewBody(team_alias=f"e2e-team-{marker}", organization_id=organization_id) + ) + self.resources.defer( + lambda: unwrap( + self.bootstrap.proxy.transport.post( + "/team/delete", + headers=self.bootstrap.proxy.management_headers(), + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + ) + ) + self.bootstrap.delete_team_member(team_id, self.bootstrap.user_info().user_id) + group_id: Final = self.idp.create_group(team_id) + self.resources.defer(lambda: self.idp.with_strict_cleanup().delete_group(group_id)) + return Tenant(organization_id=organization_id, team_id=team_id, group_id=group_id) + + def create( + self, role: ActorRole, *, tenants: tuple[Tenant, ...] = (), profile: ActorProfile = "database_role" + ) -> Actor: + if role in ("team_admin", "team_member", "organization_admin") and not tenants: + raise ValueError("A membership actor requires a tenant") + identity: Final = self.idp.with_strict_cleanup().provision_user( + marker=unique_marker(), + groups=tuple(tenant.team_id for tenant in tenants) if profile == "group_scoped" else (), + group_ids=tuple(tenant.group_id for tenant in tenants) if profile == "group_scoped" else (), + defer=self.resources.defer, + ) + global_role: Final[UserRole] = ( + role + if role in ("proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer") + else "internal_user" + ) + self.bootstrap.create_user( + UserNewBody( + user_id=identity.user_id, + user_email=f"{identity.username}@example.com", + user_role=global_role, + auto_create_key=False, + ) + ) + self.resources.defer(lambda: self.bootstrap.delete_user_strict(identity.user_id)) + for tenant in tenants: + unwrap( + self.bootstrap.proxy.transport.post( + "/organization/member_add", + headers=self.bootstrap.proxy.management_headers(), + json=OrgMemberAddBody( + organization_id=tenant.organization_id, + member=OrgMemberEntry( + user_id=identity.user_id, + role="org_admin" if role == "organization_admin" else "internal_user", + ), + ), + response_type=NoBody, + ) + ) + unwrap( + self.bootstrap.proxy.transport.post( + "/team/member_add", + headers=self.bootstrap.proxy.management_headers(), + json=TeamMemberAddBody( + team_id=tenant.team_id, + member=TeamMemberEntry( + user_id=identity.user_id, + role="admin" if role == "team_admin" else "user", + ), + ), + response_type=NoBody, + ) + ) + return Actor(identity=identity, role=role, global_role=global_role, profile=profile, tenants=tenants) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index e17b92a13ed..a0243e868e2 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -7,7 +7,8 @@ llm-only key hitting a management route). from __future__ import annotations import time -from dataclasses import dataclass +import warnings +from dataclasses import dataclass, field, replace import jwt from e2e_config import MASTER_KEY @@ -20,6 +21,7 @@ from e2e_http import ( StreamingResponse, Success, UnknownApiError, + retry_attempts, unwrap, ) from models import ( @@ -81,7 +83,7 @@ from models import ( UserNewResponse, UserUpdateBody, ) -from proxy_client import ProxyClient +from proxy_client import Caller, ProxyClient MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" @@ -98,7 +100,7 @@ class DashboardSession: its bearer on every subsequent call, the claims it renders the signed-in user from, and where it lands the browser.""" - session_key: str + session_key: str = field(repr=False) claims: UiSessionClaims redirect_url: str @@ -106,7 +108,10 @@ class DashboardSession: @dataclass(frozen=True, slots=True) class ManagementClient: proxy: ProxyClient - master_key: str + master_key: str = field(repr=False) + + def with_caller(self, caller: Caller) -> ManagementClient: + return replace(self, proxy=self.proxy.with_caller(caller)) def llm_only_key(self) -> str: return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) @@ -117,7 +122,7 @@ class ManagementClient: dashboard creates it under the session key their sign-in minted). Returns the outcome rather than unwrapping it, so a caller can poll a route that is only transiently refusing.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) return self.proxy.transport.post( "/key/generate", headers=headers, @@ -131,9 +136,9 @@ class ManagementClient: sign-in minted, never the master key). Returns the outcome rather than unwrapping it, so a caller can poll a route that is only transiently refusing; `update_key_models` is the unwrapping shorthand.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) last: Result[NoBody] = NetworkError(message="/key/update was never attempted") - for attempt in range(_KEY_WRITE_ATTEMPTS): + for attempt in range(retry_attempts(_KEY_WRITE_ATTEMPTS)): last = self.proxy.transport.post( "/key/update", headers=headers, @@ -144,6 +149,7 @@ class ManagementClient: case UnknownApiError(body=error_body) if any( marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS ): + warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2) time.sleep(0.5 * (attempt + 1)) continue case _: @@ -153,10 +159,10 @@ class ManagementClient: def update_key_models(self, key: str, models: list[str]) -> None: _ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models))) - def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]: + def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]: return self.proxy.transport.get( "/key/info", - headers=self.proxy.transport.bearer(caller_key), + headers=self.proxy.management_headers(caller_key), params=KeyInfoParams(key=key), response_type=KeyInfoResponse, ) @@ -167,7 +173,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/key/delete", - headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key), + headers=self.proxy.management_headers(caller_key), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -179,7 +185,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/model/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -190,7 +196,7 @@ class ManagementClient: Connection button, probing the live provider with the supplied params.""" return self.proxy.transport.post( "/health/test_connection", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=ConnectionTestResponse, timeout=120.0, @@ -200,7 +206,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/key/block", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyBlockBody(key=key), response_type=NoBody, ) @@ -209,7 +215,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/key/regenerate", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyRegenerateBody(key=key, grace_period=grace_period), response_type=KeyGenerateResponse, ) @@ -219,7 +225,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( f"/key/{key}/reset_spend", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyResetSpendBody(reset_to=reset_to), response_type=KeyResetSpendResponse, ) @@ -228,7 +234,7 @@ class ManagementClient: def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]: """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is who is asking: the master key by default, or a virtual key.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) return self.proxy.transport.get( "/key/list", headers=headers, @@ -266,7 +272,7 @@ class ManagementClient: team_id = unwrap( self.proxy.transport.post( "/team/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=TeamNewResponse, ) @@ -276,10 +282,10 @@ class ManagementClient: def update_team(self, body: TeamUpdateBody) -> None: last: Result[NoBody] | None = None - for attempt in range(5): + for attempt in range(retry_attempts(5)): last = self.proxy.transport.post( "/team/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -289,6 +295,7 @@ class ManagementClient: case UnknownApiError(body=body_text) if ( "connecting to redis" in body_text.lower() or "name resolution" in body_text.lower() ): + warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2) time.sleep(0.5 * (attempt + 1)) continue case _: @@ -299,7 +306,7 @@ class ManagementClient: def delete_team(self, team_id: str) -> None: _ = self.proxy.transport.post( "/team/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -308,7 +315,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/team/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -320,7 +327,7 @@ class ManagementClient: for entry in unwrap( self.proxy.transport.get( "/team/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=NoBody(), response_type=TeamListResponse, ) @@ -328,14 +335,16 @@ class ManagementClient: ) def team_info_status(self, team_id: str) -> ProbeResult: - return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) + return self.proxy.transport.probe( + "/team/info", params=TeamInfoParams(team_id=team_id), headers=self.proxy.management_headers() + ) def _wait_for_team(self, team_id: str) -> None: last: Result[TeamInfoResponse] | None = None - for _ in range(_TEAM_READY_ATTEMPTS): + for _ in range(retry_attempts(_TEAM_READY_ATTEMPTS)): last = self.proxy.transport.get( "/team/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -343,25 +352,29 @@ class ManagementClient: case Success(): return case _: + warnings.warn("Repeating team read while the team becomes available", RuntimeWarning, stacklevel=2) time.sleep(_TEAM_READY_SLEEP_SECONDS) assert last is not None raise AssertionError(last) def add_team_member(self, team_id: str, user_id: str) -> None: last: Result[NoBody] | None = None - for attempt in range(_TEAM_READY_ATTEMPTS): + for attempt in range(retry_attempts(_TEAM_READY_ATTEMPTS)): last = self.proxy.transport.post( "/team/member_add", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)), response_type=NoBody, ) match last: case Success(): return - case UnknownApiError(body=body) if ( - "doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS + case UnknownApiError(body=body) if "doesn't exist" in body and attempt + 1 < retry_attempts( + _TEAM_READY_ATTEMPTS ): + warnings.warn( + "Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2 + ) time.sleep(_TEAM_READY_SLEEP_SECONDS) continue case _: @@ -373,7 +386,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/team/member_delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id), response_type=NoBody, ) @@ -383,7 +396,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/user/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=UserNewResponse, ) @@ -393,7 +406,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/customer/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=CustomerNewBody(user_id=user_id), response_type=CustomerResponse, ) @@ -404,7 +417,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/customer/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=CustomerInfoParams(end_user_id=end_user_id), response_type=CustomerResponse, ) @@ -413,7 +426,7 @@ class ManagementClient: def delete_customer(self, user_id: str) -> None: _ = self.proxy.transport.post( "/customer/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=CustomerDeleteBody(user_ids=[user_id]), response_type=NoBody, ) @@ -422,7 +435,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/user/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -431,7 +444,7 @@ class ManagementClient: def delete_user(self, user_id: str) -> None: _ = self.proxy.transport.post( "/user/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=NoBody, ) @@ -442,17 +455,17 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/user/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=UserDeleteResponse, ) ) - def user_info(self, user_id: str) -> UserInfoResponse: + def user_info(self, user_id: str | None = None) -> UserInfoResponse: return unwrap( self.proxy.transport.get( "/user/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserInfoParams(user_id=user_id), response_type=UserInfoResponse, ) @@ -462,7 +475,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/user/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -472,7 +485,7 @@ class ManagementClient: listing = unwrap( self.proxy.transport.get( "/user/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -483,7 +496,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/organization/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=OrgNewResponse, ) @@ -493,7 +506,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.patch( "/organization/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -502,7 +515,7 @@ class ManagementClient: def delete_org(self, organization_id: str) -> None: _ = self.proxy.transport.delete( "/organization/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=OrgDeleteBody(organization_ids=[organization_id]), response_type=NoBody, ) @@ -511,19 +524,24 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/organization/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=OrgInfoParams(organization_id=organization_id), response_type=OrgInfoResponse, ) ) def org_info_status(self, organization_id: str) -> ProbeResult: - return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id)) + return self.proxy.transport.probe( + "/organization/info", + params=OrgInfoParams(organization_id=organization_id), + headers=self.proxy.management_headers(), + ) + def create_tag(self, body: TagNewBody) -> None: _ = unwrap( self.proxy.transport.post( "/tag/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -532,7 +550,7 @@ class ManagementClient: def delete_tag(self, name: str) -> None: _ = self.proxy.transport.post( "/tag/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TagDeleteBody(name=name), response_type=NoBody, ) @@ -542,7 +560,7 @@ class ManagementClient: unwrap( self.proxy.transport.get( "/tag/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=NoBody(), response_type=TagListResponse, ) @@ -553,7 +571,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/v1/mcp/server", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=McpServerRow, ) @@ -565,7 +583,7 @@ class ManagementClient: return unwrap( self.proxy.transport.put( "/v1/mcp/server", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=McpServerRow, ) @@ -576,7 +594,7 @@ class ManagementClient: unwrap it while a deferred teardown can ignore an already-deleted server.""" return self.proxy.transport.delete( f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=NoBody(), response_type=NoBody, ) diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py index 22306da8eb8..0d23f954158 100644 --- a/tests/e2e/management/test_jwt_management_e2e.py +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -2,60 +2,249 @@ from __future__ import annotations -from typing import Final +from typing import Final, Literal import pytest -from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker from e2e_http import UnauthorizedError, UnknownApiError, unwrap -from idp import ADMIN_CLIENT_ID, Identity, Keycloak +from idp import ADMIN_CLIENT_ID, Identity, Keycloak, token_claims from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory, ActorRole from management_client import ManagementClient -from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody +from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserInfoParams, UserInfoResponse, UserNewBody +from proxy_client import Caller pytestmark = pytest.mark.e2e class TestJwtManagement: + @pytest.mark.parametrize( + "role", + ( + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", + ), + ) + @pytest.mark.covers("mgmt.user.jwt.database_roles") + def test_actor_subject_and_database_role(self, actor_factory: ActorFactory, role: ActorRole) -> None: + tenants: Final = ( + (actor_factory.tenant(),) if role in ("organization_admin", "team_admin", "team_member") else () + ) + actor: Final = actor_factory.create(role, tenants=tenants) + caller: Final = actor.mint_caller(actor_factory.idp) + claims: Final = token_claims(caller.credential) + assert claims.sub == actor.identity.user_id + assert claims.iss == actor_factory.idp.issuer + assert claims.aud == "litellm-e2e" or "litellm-e2e" in claims.aud + assert actor.identity.groups == () + assert ("litellm_proxy_admin" in claims.scope.split()) == (role == "proxy_admin") + stored: Final = actor_factory.bootstrap.user_info(actor.identity.user_id) + assert stored.user_id == actor.identity.user_id + assert stored.user_info.user_role == actor.global_role + bound: Final = actor_factory.bootstrap.with_caller(caller) + own: Final = unwrap( + bound.proxy.transport.get( + "/user/info", + headers=bound.proxy.management_headers(), + params=UserInfoParams(), + response_type=UserInfoResponse, + ) + ) + assert own.user_id == actor.identity.user_id + assert own.user_info.user_role == actor.global_role + for tenant in tenants: + info = actor_factory.bootstrap.team_info(tenant.team_id) + assert info.organization_id == tenant.organization_id + assert {(member.user_id, member.role) for member in info.members_with_roles} == { + (actor.identity.user_id, "admin" if role == "team_admin" else "user") + } + assert { + (member.user_id, member.user_role) + for member in actor_factory.bootstrap.org_info(tenant.organization_id).members + } == {(actor.identity.user_id, "org_admin" if role == "organization_admin" else "internal_user")} + + @pytest.mark.covers("mgmt.key.jwt.viewer_denied") + def test_admin_viewer_reads_but_cannot_update(self, actor_factory: ActorFactory) -> None: + actor: Final = actor_factory.create("proxy_admin_viewer") + viewer: Final = actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) + alias: Final = f"e2e-viewer-{unique_marker()}" + key: Final = actor_factory.key().key + unwrap(actor_factory.bootstrap.update_key(KeyUpdateBody(key=key, key_alias=alias))) + assert viewer.proxy.key_info(key).key_alias == alias + denied: Final = viewer.update_key(KeyUpdateBody(key=key, key_alias="forbidden")) + assert isinstance(denied, UnknownApiError) and denied.status_code == 403, f"viewer write was accepted: {denied}" + assert "proxy_admin_viewer" in denied.body and "/key/update" in denied.body + assert actor_factory.bootstrap.proxy.key_info(key).key_alias == alias + + @pytest.mark.covers("mgmt.user.oidc.identity_mapping") + def test_oidc_browser_profile_identity_mapping(self, actor_factory: ActorFactory) -> None: + actor: Final = actor_factory.create("internal_user") + idp: Final = actor_factory.idp.with_strict_cleanup() + discovery: Final = idp.discovery() + assert discovery.issuer == idp.issuer + assert discovery.jwks_uri == idp.jwks_url + callback: Final = f"{PROXY_BASE_URL}/sso/callback" + browser: Final = idp.browser_client(callback_url=callback, defer=actor_factory.resources.defer) + token: Final = idp.browser_token(actor.identity, browser) + assert token_claims(token).sub == actor.identity.user_id + userinfo: Final = idp.userinfo(token) + assert userinfo.sub == actor.identity.user_id + assert userinfo.email == f"{actor.identity.username}@example.com" + assert browser.environment(discovery)["GENERIC_USER_ID_ATTRIBUTE"] == "sub" + @pytest.mark.covers("mgmt.key.jwt.lifecycle") + @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key")) def test_admin_creates_reads_updates_clears_and_deletes_a_key( - self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager + self, + client: ManagementClient, + idp: Keycloak, + jwt_identity: Identity, + resources: ResourceManager, + actor_factory: ActorFactory, + credential_kind: Literal["direct_jwt", "virtual_key"], ) -> None: - admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + actor: Final = actor_factory.create("proxy_admin") + virtual_key: Final = ( + actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None + ) + admin: Final = ( + virtual_key if virtual_key is not None else idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + ) + bound: Final = client.with_caller(Caller(credential=admin, kind=credential_kind, role="proxy_admin")) alias: Final = f"e2e-jwt-key-{unique_marker()}" created: Final = unwrap( - client.generate_key( + bound.generate_key( KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), - caller_key=admin, ) ) resources.defer(lambda: client.proxy.delete_key(created.key)) - original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + original: Final = unwrap(bound.key_info_as(created.key)).info assert original.key_alias == alias and original.team_id == jwt_identity.group assert original.models == [CHEAP_OPENAI_MODEL] updated_alias: Final = f"{alias}-updated" - unwrap( - client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin) - ) - updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + unwrap(bound.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120))) + updated: Final = unwrap(bound.key_info_as(created.key)).info assert updated.key_alias == updated_alias and updated.rpm_limit == 120 assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction" - unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin)) - cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + unwrap(bound.update_key(KeyUpdateBody(key=created.key, models=[]))) + cleared: Final = unwrap(bound.key_info_as(created.key)).info assert cleared.models == [] and cleared.rpm_limit == 120 - assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1 - client.delete_key_strict(created.key, caller_key=admin) - assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0 + assert unwrap(bound.key_list(updated_alias)).total_count == 1 + bound.delete_key_strict(created.key) + assert unwrap(bound.key_list(updated_alias)).total_count == 0 + + @pytest.mark.covers("mgmt.team.jwt.tenant_isolation") + def test_two_actor_sets_keep_tenants_and_keys_isolated(self, actor_factory: ActorFactory) -> None: + first: Final = actor_factory.tenant() + second: Final = actor_factory.tenant() + assert first.organization_id != second.organization_id and first.team_id != second.team_id + actors: Final = tuple( + actor_factory.create("team_member", tenants=(tenant,), profile="group_scoped") for tenant in (first, second) + ) + assert actors[0].identity.user_id != actors[1].identity.user_id + callers: Final = tuple( + actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) for actor in actors + ) + keys: Final = tuple(actor_factory.key(tenant) for tenant in (first, second)) + assert keys[0].key != keys[1].key + assert callers[0].proxy.key_info(keys[0].key).team_id == first.team_id + assert callers[1].proxy.key_info(keys[1].key).team_id == second.team_id + for caller, other_key in ((callers[0], keys[1].key), (callers[1], keys[0].key)): + hidden = caller.key_info_as(other_key) + assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403 + assert tuple(actor.identity.groups for actor in actors) == ((first.team_id,), (second.team_id,)) + + @pytest.mark.covers("mgmt.team.jwt.multiple_memberships") + def test_multi_group_actor_keeps_exact_memberships(self, actor_factory: ActorFactory) -> None: + tenants: Final = (actor_factory.tenant(), actor_factory.tenant()) + actor: Final = actor_factory.create("team_member", tenants=tenants, profile="group_scoped") + claims: Final = token_claims(actor.mint_caller(actor_factory.idp).credential) + assert set(claims.groups) == {tenant.team_id for tenant in tenants} + assert "litellm_proxy_admin" not in claims.scope.split() + assert actor.identity.groups == tuple(tenant.team_id for tenant in tenants) + for tenant in tenants: + assert { + (entry.user_id, entry.role) + for entry in actor_factory.bootstrap.team_info(tenant.team_id).members_with_roles + } == {(actor.identity.user_id, "user")} + + @pytest.mark.covers("mgmt.user.jwt.cleanup") + def test_successful_actor_cleanup_removes_owned_state(self, actor_factory: ActorFactory) -> None: + resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True) + factory: Final = ActorFactory(bootstrap=actor_factory.bootstrap, idp=actor_factory.idp, resources=resources) + try: + tenant: Final = factory.tenant() + actor: Final = factory.create("team_member", tenants=(tenant,), profile="group_scoped") + key: Final = factory.key(tenant) + alias: Final = factory.bootstrap.proxy.key_info(key.key).key_alias + assert alias is not None + finally: + resources.teardown() + assert factory.bootstrap.user_count(actor.identity.user_id) == 0 + assert factory.bootstrap.key_alias_count(alias) == 0 + assert factory.bootstrap.team_info_status(tenant.team_id).status_code == 404 + assert factory.bootstrap.org_info_status(tenant.organization_id).status_code == 404 + factory.idp.assert_absent("users", actor.identity.user_id) + factory.idp.assert_absent("groups", tenant.group_id) + + @pytest.mark.parametrize("stage", ("group", "user")) + @pytest.mark.covers("mgmt.user.jwt.partial_cleanup") + def test_partial_setup_removes_previously_created_identities( + self, + actor_factory: ActorFactory, + stage: Literal["group", "user"], + ) -> None: + idp: Final = actor_factory.idp.with_strict_cleanup() + resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True) + marker: Final = unique_marker() + group_id: Final = idp.create_group(f"e2e-partial-{marker}") + resources.defer(lambda: idp.delete_group(group_id)) + try: + identity: Final = ( + idp.provision_user( + marker=marker, + groups=(f"e2e-partial-{marker}",), + group_ids=(group_id,), + defer=resources.defer, + ) + if stage == "user" + else None + ) + if identity is None: + with pytest.raises(pytest.fail.Exception, match="HTTP 409"): + idp.create_group(f"e2e-partial-{marker}") + else: + with pytest.raises(pytest.fail.Exception, match="HTTP 409"): + idp.create_user( + username=identity.username, + email=f"{identity.username}@example.com", + password=identity.password, + groups=identity.groups, + ) + finally: + resources.teardown() + idp.assert_absent("groups", group_id) + if identity is not None: + idp.assert_absent("users", identity.user_id) @pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied") def test_member_cannot_write_and_another_team_cannot_read_the_key( self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager ) -> None: admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + bound: Final = client.with_caller(Caller(credential=admin, kind="direct_jwt", role="proxy_admin")) member: Final = idp.access_token(jwt_identity) + member_client: Final = client.with_caller(Caller(credential=member, kind="direct_jwt", role="team_member")) alias: Final = f"e2e-jwt-owned-{unique_marker()}" created: Final = unwrap( client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin) @@ -63,14 +252,14 @@ class TestJwtManagement: resources.defer(lambda: client.proxy.delete_key(created.key)) client.add_team_member(jwt_identity.group, jwt_identity.user_id) - assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias + assert unwrap(member_client.key_info_as(created.key)).info.key_alias == alias - refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member) + refused: Final = member_client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden")) assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}" assert "does not have permissions for endpoint" in refused.body.lower(), ( f"expected a permission denial: {refused}" ) - assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias + assert unwrap(bound.key_info_as(created.key)).info.key_alias == alias marker: Final = unique_marker() outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer) @@ -88,4 +277,4 @@ class TestJwtManagement: assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, ( f"another team must not read this key: {hidden}" ) - assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group + assert unwrap(bound.key_info_as(created.key)).info.team_id == jwt_identity.group diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 3cab0334dea..6fca1268ebc 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1091,13 +1091,13 @@ class UiLoginBody(BaseModel): class UiLoginResponse(BaseModel): - token: str + token: str = Field(repr=False) redirect_url: str class UiSessionClaims(BaseModel): user_id: str - key: str + key: str = Field(repr=False) user_role: str login_method: Literal["sso", "username_password"] exp: int @@ -1135,6 +1135,7 @@ class TeamInfoParams(BaseModel): class TeamData(BaseModel): + organization_id: str | None = None team_alias: str | None = None models: list[str] = [] members_with_roles: list[TeamMemberEntry] = [] @@ -1175,6 +1176,7 @@ class UserNewBody(BaseModel): user_email: str user_role: UserRole user_id: str | None = None + auto_create_key: bool | None = None class UserNewResponse(BaseModel): @@ -1187,7 +1189,7 @@ class UserUpdateBody(BaseModel): class UserInfoParams(BaseModel): - user_id: str + user_id: str | None = None class UserData(BaseModel): @@ -1240,16 +1242,36 @@ class OrgInfoParams(BaseModel): organization_id: str +class OrgMembership(BaseModel): + user_id: str + user_role: str + + class OrgInfoResponse(BaseModel): organization_id: str organization_alias: str | None = None models: list[str] = [] + members: tuple[OrgMembership, ...] = () + + +class OrgMemberEntry(BaseModel): + user_id: str + role: Literal["org_admin", "internal_user"] + + +class OrgMemberAddBody(BaseModel): + organization_id: str + member: OrgMemberEntry class OrgDeleteBody(BaseModel): organization_ids: list[str] +class OrgDeleteResponse(RootModel[tuple[OrgInfoResponse, ...]]): + pass + + # ---------- tags (management) ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1fe2ec905ef..48a6110dc0b 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -11,14 +11,23 @@ from __future__ import annotations import time import warnings from collections.abc import Callable, Mapping -from dataclasses import dataclass -from functools import reduce +from dataclasses import dataclass, field, replace from datetime import datetime +from functools import reduce from types import MappingProxyType -from typing import Final - -from pydantic import BaseModel +from typing import Final, Literal +from e2e_config import ( + CONTROL_PLANE_BASE_URL, + MASTER_KEY, + POLL_INTERVAL, + POLL_TIMEOUT, + PROXY_BASE_URL, + PROXY_REPLICA_URLS, + REQUEST_TIMEOUT, + SLOW_PROVIDER_TIMEOUT_SECONDS, + settle_propagation, +) from e2e_http import ( AnthropicHeaders, AuthHeaders, @@ -55,6 +64,7 @@ from models import ( KeyInfoParams, KeyInfoResponse, LiteLLMParamsBody, + MemorySummaryResponse, ModelDeleteBody, ModelInfoBody, ModelInfoEntry, @@ -63,7 +73,6 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListParams, - MemorySummaryResponse, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -76,23 +85,13 @@ from models import ( TeamDeleteBody, TeamNewBody, TeamNewResponse, - UserDeleteBody, - UserDeleteResponse, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, + UserDeleteBody, + UserDeleteResponse, ) -from e2e_config import ( - CONTROL_PLANE_BASE_URL, - MASTER_KEY, - POLL_INTERVAL, - POLL_TIMEOUT, - PROXY_BASE_URL, - PROXY_REPLICA_URLS, - REQUEST_TIMEOUT, - SLOW_PROVIDER_TIMEOUT_SECONDS, - settle_propagation, -) +from pydantic import BaseModel from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -421,11 +420,23 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re ) +CredentialKind = Literal["master", "direct_jwt", "virtual_key", "dashboard_session"] + + +@dataclass(frozen=True, slots=True) +class Caller: + credential: str = field(repr=False) + kind: CredentialKind + role: str + tenant: str | None = None + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport replicas: Mapping[str, Transport] control_replicas: Mapping[str, Transport] + caller: Caller | None = None poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -433,13 +444,24 @@ class ProxyClient: model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT + def with_caller(self, caller: Caller) -> ProxyClient: + return replace(self, caller=caller) + + def management_headers(self, caller_key: str | None = None, *, transport: Transport | None = None) -> AuthHeaders: + selected: Final = self.transport if transport is None else transport + if caller_key is not None: + return selected.bearer(caller_key) + if self.caller is not None: + return selected.bearer(self.caller.credential) + return selected.master + # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- def generate_key(self, body: KeyGenerateBody) -> str: return unwrap( self.transport.post( "/key/generate", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=KeyGenerateResponse, ) @@ -448,7 +470,7 @@ class ProxyClient: def delete_key(self, key: str) -> None: _ = self.transport.post( "/key/delete", - headers=self.transport.master, + headers=self.management_headers(), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -458,7 +480,7 @@ class ProxyClient: return _ = self.transport.post( "/customer/delete", - headers=self.transport.master, + headers=self.management_headers(), json=CustomerDeleteBody(user_ids=user_ids), response_type=NoBody, ) @@ -467,7 +489,7 @@ class ProxyClient: return unwrap( self.transport.get( "/key/info", - headers=self.transport.master, + headers=self.management_headers(), params=KeyInfoParams(key=key), response_type=KeyInfoResponse, ) @@ -477,7 +499,7 @@ class ProxyClient: return { url: transport.get( "/debug/memory/summary", - headers=transport.master, + headers=self.management_headers(transport=transport), params=NoBody(), response_type=MemorySummaryResponse, ) @@ -524,11 +546,12 @@ class ProxyClient: {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} ) - @staticmethod def _body_poller[R: BaseModel]( - transport: Transport, path: str, params: BaseModel, response_type: type[R] + self, transport: Transport, path: str, params: BaseModel, response_type: type[R] ) -> Poller[Result[R]]: - return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + return lambda: transport.get( + path, headers=self.management_headers(transport=transport), params=params, response_type=response_type + ) def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it @@ -536,7 +559,7 @@ class ProxyClient: return unwrap( self.transport.get( "/model/info", - headers=self.transport.master, + headers=self.management_headers(), params=NoBody(), response_type=ModelInfoResponse, ) @@ -546,7 +569,7 @@ class ProxyClient: return unwrap( self.transport.get( "/public/litellm_model_cost_map", - headers=self.transport.master, + headers=self.management_headers(), params=NoBody(), response_type=CostMap, ) @@ -607,7 +630,7 @@ class ProxyClient: model_id = unwrap( self.transport.post( "/model/new", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ModelNewResponse, ) @@ -623,7 +646,7 @@ class ProxyClient: def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" - headers: Final = self.transport.master if listed_for is None else self.transport.bearer(listed_for) + headers: Final = self.management_headers(listed_for) outcome: Final = await_servable_everywhere( {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, @@ -666,7 +689,7 @@ class ProxyClient: unwrap( self.transport.post( "/model/update", - headers=self.transport.master, + headers=self.management_headers(), json=ModelUpdateBody( litellm_params=litellm_params, model_info=ModelInfoBody(id=model_id), @@ -678,7 +701,7 @@ class ProxyClient: def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", - headers=self.transport.master, + headers=self.management_headers(), json=ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -747,11 +770,10 @@ class ProxyClient: f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" ) - @staticmethod - def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + def _reader[R: BaseModel](self, transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: return lambda request_timeout: transport.get( path, - headers=transport.master, + headers=self.management_headers(transport=transport), params=NoBody(), response_type=response_type, timeout=request_timeout, @@ -763,7 +785,7 @@ class ProxyClient: return unwrap( self.transport.post( "/v1/mcp/toolset", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ToolsetRow, ) @@ -775,7 +797,7 @@ class ProxyClient: return unwrap( self.transport.put( "/v1/mcp/toolset", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ToolsetRow, ) @@ -786,7 +808,7 @@ class ProxyClient: can unwrap it while a deferred teardown can ignore an already-deleted row.""" return self.transport.delete( f"/v1/mcp/toolset/{toolset_id}", - headers=self.transport.master, + headers=self.management_headers(), json=NoBody(), response_type=NoBody, ) @@ -795,7 +817,7 @@ class ProxyClient: unwrap( self.transport.post( "/credentials", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=CredentialCreateResponse, ) @@ -804,7 +826,7 @@ class ProxyClient: def delete_credential(self, credential_name: str) -> None: result = self.transport.delete( f"/credentials/{credential_name}", - headers=self.transport.master, + headers=self.management_headers(), json=NoBody(), response_type=NoBody, ) @@ -815,7 +837,7 @@ class ProxyClient: return unwrap( self.transport.post( "/team/new", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=TeamNewResponse, ) @@ -824,7 +846,7 @@ class ProxyClient: def delete_team(self, team_id: str) -> None: result = self.transport.post( "/team/delete", - headers=self.transport.master, + headers=self.management_headers(), json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -836,7 +858,7 @@ class ProxyClient: a user the proxy only upserts after a successful auth.""" result = self.transport.post( "/user/delete", - headers=self.transport.master, + headers=self.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=UserDeleteResponse, ) @@ -909,7 +931,7 @@ class ProxyClient: def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]: result = self.transport.get( "/spend/logs", - headers=self.transport.master, + headers=self.management_headers(), params=params, response_type=SpendLogs, ) @@ -924,7 +946,7 @@ class ProxyClient: return unwrap( self.transport.get( "/spend/logs/v2", - headers=self.transport.master, + headers=self.management_headers(), params=SpendLogsPageParams( start_date=start.strftime("%Y-%m-%d %H:%M:%S"), end_date=end.strftime("%Y-%m-%d %H:%M:%S"), @@ -977,7 +999,7 @@ class ProxyClient: # ---- route probe ---------------------------------------------------- def probe(self, path: str, *, params: NoBody) -> ProbeResult: - return self.transport.probe(path, params=params) + return self.transport.probe(path, params=params, headers=self.management_headers()) def build_proxy_client( diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 81cd6c8d3d1..7201da84924 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -29,6 +29,7 @@ from e2e_http import ( request_with_retry, streaming_outcome, wire_body, + without_retries, ) from pydantic import BaseModel, TypeAdapter @@ -56,6 +57,15 @@ def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse] class TestTransientRetryPolicy: + def test_qualification_disables_retries_and_restores_the_default(self) -> None: + responses: Final = (FakeResponse(529), FakeResponse(200)) + sleep: Final = SleepRecorder() + with without_retries(): + assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[0] + assert sleep.delays == () + assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[1] + assert sleep.delays == (0.5,) + def test_transient_set_is_only_statuses_the_proxy_cannot_emit(self) -> None: assert TRANSIENT_STATUSES == frozenset({529}) assert 429 not in TRANSIENT_STATUSES diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py index 33a09a0f13a..8a92dfa527c 100644 --- a/tests/e2e/test_idp.py +++ b/tests/e2e/test_idp.py @@ -4,12 +4,19 @@ these carry no `e2e` marker and run everywhere.""" from __future__ import annotations +import os +import signal +import subprocess +import sys +import time +from builtins import ExceptionGroup from collections.abc import Callable, Generator from contextlib import ExitStack, contextmanager from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path from queue import SimpleQueue from threading import Thread -from typing import Final +from typing import Final, Literal import pytest from e2e_http import ExternalWrite @@ -18,6 +25,8 @@ from idp import ( KEYCLOAK_ADMIN_USER_ENV, KEYCLOAK_REALM_ENV, KEYCLOAK_URL_ENV, + BrowserClientBody, + Discovery, Keycloak, PasswordCredential, UserCreateBody, @@ -60,24 +69,48 @@ def _idp_server( ) -> Generator[tuple[Keycloak, SimpleQueue[str]]]: """Exercise provisioning failures through the same HTTP transport as live tests.""" deletions: SimpleQueue[str] = SimpleQueue() + clients: SimpleQueue[BrowserClientBody] = SimpleQueue() class Handler(BaseHTTPRequestHandler): def log_message(self, format: str, *args: object) -> None: pass def do_POST(self) -> None: - self.rfile.read(int(self.headers.get("Content-Length", "0"))) + body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) if self.path.endswith("/token"): self.send_response(admin_status) self.end_headers() self.wfile.write(b'{"access_token":"synthetic-harness-token"}') else: + if self.path.endswith("/clients"): + clients.put(BrowserClientBody.model_validate_json(body)) self.send_response(user_status if self.path.endswith("/users") else 201) self.send_header("Location", f"{self.path}/resource-1") self.end_headers() if user_status != 201 and self.path.endswith("/users"): self.wfile.write(b"injected create failure") + def do_GET(self) -> None: + self.send_response(200) + self.end_headers() + if "/clients/" in self.path: + client: Final = clients.get_nowait() + clients.put(client) + self.wfile.write(client.model_dump_json(by_alias=True).encode()) + else: + issuer: Final = f"http://127.0.0.1:{server.server_port}/realms/test" + self.wfile.write( + Discovery( + issuer=issuer, + authorization_endpoint=f"{issuer}/auth", + token_endpoint=f"{issuer}/token", + userinfo_endpoint=f"{issuer}/userinfo", + jwks_uri=f"{issuer}/certs", + ) + .model_dump_json() + .encode() + ) + def do_DELETE(self) -> None: deletions.put(self.path) self.send_response(delete_status) @@ -115,6 +148,54 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No assert deletions.empty() +@pytest.mark.parametrize("exit_mode", ("normal", "parent", "group")) +def test_oidc_launcher_removes_client_on_exit_and_termination( + tmp_path: Path, exit_mode: Literal["normal", "parent", "group"] +) -> None: + ready: Final = tmp_path / "ready" + child_command: Final = ( + "import os,time; from pathlib import Path; " + 'assert os.environ["GENERIC_CLIENT_SECRET"]; ' + 'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; ' + f"Path({str(ready)!r}).touch(); " + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") + ) + with _idp_server() as (idp, deletions): + with subprocess.Popen( + [ + sys.executable, + str(Path(__file__).with_name("idp.py")), + "http://127.0.0.1:9999", + sys.executable, + "-c", + child_command, + ], + env={ + **os.environ, + KEYCLOAK_URL_ENV: idp.base_url, + KEYCLOAK_REALM_ENV: idp.realm, + KEYCLOAK_ADMIN_USER_ENV: idp.admin_username, + KEYCLOAK_ADMIN_PASSWORD_ENV: idp.admin_password, + }, + start_new_session=True, + ) as process: + try: + deadline: Final = time.monotonic() + 15 + while not ready.exists() and time.monotonic() < deadline and process.poll() is None: + time.sleep(0.05) + assert ready.exists(), "OIDC child did not start" + if exit_mode == "parent": + process.terminate() + elif exit_mode == "group": + os.killpg(process.pid, signal.SIGTERM) + assert process.wait(timeout=10) == (7 if exit_mode == "normal" else 143) + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=5) + assert deletions.get(timeout=5) == "/admin/realms/test/clients/resource-1" + assert deletions.empty() + + def test_successful_provisioning_cleans_up_user_before_group() -> None: with _idp_server() as (idp, deletions): with ExitStack() as cleanup: @@ -134,6 +215,43 @@ def test_cleanup_failure_is_visible() -> None: idp.delete_group("group") +def test_strict_cleanup_reports_each_failure_and_continues() -> None: + from lifecycle import ResourceManager + from proxy_client import build_proxy_client + + with _idp_server(delete_status=500) as (idp, deletions): + resources: Final = ResourceManager(client=build_proxy_client(), strict_cleanup=True) + strict: Final = idp.with_strict_cleanup() + resources.defer(lambda: strict.delete_group("group")) + resources.defer(lambda: strict.delete_user("user")) + with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as error: + resources.teardown() + assert len(error.value.exceptions) == 2 + assert deletions.get_nowait() == "/admin/realms/test/users/user" + assert deletions.get_nowait() == "/admin/realms/test/groups/group" + + +@pytest.mark.parametrize("groups", ((), ("one",), ("one", "two"))) +def test_provisioning_records_zero_one_or_multiple_groups(groups: tuple[str, ...]) -> None: + with _idp_server() as (idp, deletions): + with ExitStack() as cleanup: + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + identity: Final = idp.provision_groups( + marker="memberships", + groups=groups, + defer=defer, + ) + assert identity.groups == groups + assert len(identity.group_ids) == len(groups) + assert deletions.get_nowait() == "/admin/realms/test/users/resource-1" + for _ in groups: + assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1" + assert deletions.empty() + + def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None: with _idp_server(admin_status=401) as (idp, _): cleanup: Final = ExitStack() diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 1b0133f12cb..879bd88980c 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -11,19 +11,53 @@ injected clock, so nothing here monkeypatches anything. from __future__ import annotations -from collections.abc import Iterable, Mapping +import json +from builtins import ExceptionGroup +from collections.abc import Callable, Generator, Iterable, Mapping +from contextlib import contextmanager from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain, repeat +from queue import SimpleQueue +from threading import Thread from types import MappingProxyType from typing import Final, cast import pytest from e2e_config import parse_replica_urls -from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from e2e_http import NoBody, Result, Success, without_retries +from idp import Keycloak +from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory +from management.management_client import ManagementClient +from models import ( + ConnectionTestBody, + CredentialCreateBody, + KeyGenerateBody, + KeyInfo, + KeyInfoResponse, + KeyUpdateBody, + LiteLLMParamsBody, + McpServerCreateBody, + McpServerUpdateBody, + ModelListEntry, + ModelsListResponse, + OrgNewBody, + OrgUpdateBody, + SpendLogsParams, + TagNewBody, + TeamNewBody, + TeamUpdateBody, + ToolsetCreateBody, + ToolsetUpdateBody, + UserNewBody, + UserUpdateBody, +) from proxy_client import ( - ConvergeOutcome, + Caller, Converged, + ConvergeOutcome, + CredentialKind, EverywhereConverged, ModelsPoller, NeverConvergedOn, @@ -42,6 +76,108 @@ from proxy_client import ( ) from transport import Transport + +@contextmanager +def caller_boundary( + status: int = 200, bodies: SimpleQueue[bytes] | None = None, *, delete_status: int | None = None +) -> Generator[tuple[ManagementClient, SimpleQueue[str]]]: + received: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def do_GET(self) -> None: + received.put(self.headers.get("Authorization", "")) + self.send_response(delete_status if self.path == "/key/delete" and delete_status is not None else status) + self.end_headers() + self.wfile.write( + b'{"key":"owned","info":{"key_alias":"owned"},"data":[{"id":"owned"}],"team_id":"owned","team_info":{},"model_id":"owned"}' + ) + + def do_POST(self) -> None: + body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if bodies is not None: + bodies.put(body) + self.do_GET() + + do_PATCH = do_POST + do_PUT = do_POST + do_DELETE = do_POST + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + proxy: Final = build_proxy_client( + base_url=url, control_plane_base_url=url, replica_urls=(url,), master_key="bootstrap" + ) + try: + yield ManagementClient(proxy=proxy, master_key="bootstrap"), received + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestBoundManagementCaller: + def test_actor_key_cleanup_reports_failure_and_continues(self) -> None: + with caller_boundary(delete_status=500) as (bootstrap, received), without_retries(): + resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True) + remaining: SimpleQueue[str] = SimpleQueue() + resources.defer(lambda: remaining.put("cleaned")) + factory: Final = ActorFactory( + bootstrap=bootstrap, + idp=Keycloak(base_url="http://unused.test", realm="test", admin_username="test", admin_password="test"), + resources=resources, + ) + assert factory.key().key == "owned" + with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as failure: + resources.teardown() + assert len(failure.value.exceptions) == 1 + assert remaining.get_nowait() == "cleaned" + assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap") + + @pytest.mark.parametrize("kind", ("direct_jwt", "virtual_key", "dashboard_session")) + def test_direct_delegated_and_replica_reads_keep_the_bound_caller(self, kind: CredentialKind) -> None: + with caller_boundary() as (bootstrap, received): + caller: Final = Caller(credential="synthetic-caller", kind=kind, role="internal_user", tenant="tenant-a") + bound: Final = bootstrap.with_caller(caller) + bound.update_key(KeyUpdateBody(key="owned", key_alias="updated")) + bound.proxy.key_info("owned") + bound.proxy.read_back_everywhere( + "/key/info", + params=KeyUpdateBody(key="owned"), + response_type=KeyInfoResponse, + converged=lambda result: isinstance(result, Success), + ) + bound.proxy.read_body_back_everywhere( + "/key/info", KeyInfoResponse, settled=lambda result: result.info.key_alias == "owned" + ) + assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer synthetic-caller",) * 4 + assert received.empty() + bootstrap.proxy.key_info("owned") + assert received.get_nowait() == "Bearer bootstrap" + + def test_explicit_override_wins_without_rebinding_or_changing_master(self) -> None: + with caller_boundary() as (bootstrap, received): + bound: Final = bootstrap.with_caller(Caller(credential="bound", kind="direct_jwt", role="internal_user")) + bound.update_key(KeyUpdateBody(key="owned"), caller_key="override") + bound.proxy.key_info("owned") + assert received.get_nowait() == "Bearer override" + assert received.get_nowait() == "Bearer bound" + assert bound.master_key == "bootstrap" + + def test_credentials_are_absent_from_binding_and_header_diagnostics(self) -> None: + with caller_boundary() as (bootstrap, _): + caller: Final = Caller(credential="private-value", kind="direct_jwt", role="internal_user") + bound: Final = bootstrap.with_caller(caller) + assert "private-value" not in repr(caller) + assert "private-value" not in repr(bound) + assert "private-value" not in repr(bound.proxy.management_headers()) + assert "bootstrap" not in repr(bound) + + MODEL: Final = "gpt-under-test" _NO_TRANSPORTS: Final = cast(Transport, None) TIMEOUT: Final = 10.0 @@ -275,3 +411,166 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") + + +MANAGEMENT_OPERATIONS: Final[tuple[tuple[str, Callable[[ManagementClient], object]], ...]] = ( + ("generate_key", lambda c: c.generate_key(KeyGenerateBody())), + ("llm_only_key", lambda c: c.llm_only_key()), + ("update_key", lambda c: c.update_key(KeyUpdateBody(key="owned"))), + ("update_key_models", lambda c: c.update_key_models("owned", [])), + ("key_info", lambda c: c.key_info_as("owned")), + ("delete_key_strict", lambda c: c.delete_key_strict("owned")), + ("delete_model_strict", lambda c: c.delete_model_strict("owned")), + ( + "connection_test", + lambda c: c.connection_test( + ConnectionTestBody(litellm_params=LiteLLMParamsBody(model="synthetic"), mode="chat") + ), + ), + ("block_key", lambda c: c.block_key("owned")), + ("regenerate_key", lambda c: c.regenerate_key("owned")), + ("reset_key_spend", lambda c: c.reset_key_spend("owned", 0)), + ("key_list", lambda c: c.key_list("owned")), + ("key_alias_count", lambda c: c.key_alias_count("owned")), + ("create_team", lambda c: c.create_team(TeamNewBody(team_alias="owned"))), + ("update_team", lambda c: c.update_team(TeamUpdateBody(team_id="owned", team_alias="updated"))), + ("delete_team", lambda c: c.delete_team("owned")), + ("team_info", lambda c: c.team_info("owned")), + ("team_list_ids", lambda c: c.team_list_ids()), + ("team_info_status", lambda c: c.team_info_status("owned")), + ("add_team_member", lambda c: c.add_team_member("owned", "user")), + ("delete_team_member", lambda c: c.delete_team_member("owned", "user")), + ("create_user", lambda c: c.create_user(UserNewBody(user_email="actor@example.com", user_role="internal_user"))), + ("create_customer", lambda c: c.create_customer("owned")), + ("customer_info", lambda c: c.customer_info("owned")), + ("delete_customer", lambda c: c.delete_customer("owned")), + ("update_user", lambda c: c.update_user(UserUpdateBody(user_id="owned", user_role="internal_user"))), + ("delete_user", lambda c: c.delete_user("owned")), + ("delete_user_strict", lambda c: c.delete_user_strict("owned")), + ("user_info", lambda c: c.user_info("owned")), + ("user_count", lambda c: c.user_count("owned")), + ("user_list_ids", lambda c: c.user_list_ids("owned")), + ("create_org", lambda c: c.create_org(OrgNewBody(organization_alias="owned"))), + ("update_org", lambda c: c.update_org(OrgUpdateBody(organization_id="owned", organization_alias="updated"))), + ("delete_org", lambda c: c.delete_org("owned")), + ("org_info", lambda c: c.org_info("owned")), + ("org_info_status", lambda c: c.org_info_status("owned")), + ("create_tag", lambda c: c.create_tag(TagNewBody(name="owned"))), + ("delete_tag", lambda c: c.delete_tag("owned")), + ("tag_list", lambda c: c.tag_list()), + ("create_mcp_server", lambda c: c.create_mcp_server(McpServerCreateBody(alias="owned", url="http://example.test"))), + ("update_mcp_server", lambda c: c.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))), + ("delete_mcp_server", lambda c: c.delete_mcp_server("owned")), + ("proxy.generate_key", lambda c: c.proxy.generate_key(KeyGenerateBody())), + ("proxy.delete_key", lambda c: c.proxy.delete_key("owned")), + ("proxy.delete_customers", lambda c: c.proxy.delete_customers(["owned"])), + ("proxy.key_info", lambda c: c.proxy.key_info("owned")), + ("proxy.memory_summary", lambda c: c.proxy.memory_summary_everywhere()), + ("proxy.model_info", lambda c: c.proxy.model_info()), + ("proxy.model_cost_map", lambda c: c.proxy.model_cost_map()), + ("proxy.create_model", lambda c: c.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))), + ("proxy.update_model", lambda c: c.proxy.update_model("owned", LiteLLMParamsBody(model="synthetic"))), + ("proxy.delete_model", lambda c: c.proxy.delete_model("owned")), + ("proxy.create_toolset", lambda c: c.proxy.create_toolset(ToolsetCreateBody(toolset_name="owned", tools=[]))), + ("proxy.update_toolset", lambda c: c.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))), + ("proxy.delete_toolset", lambda c: c.proxy.delete_toolset("owned")), + ( + "proxy.create_credential", + lambda c: c.proxy.create_credential(CredentialCreateBody(credential_name="owned", credential_values={})), + ), + ("proxy.delete_credential", lambda c: c.proxy.delete_credential("owned")), + ("proxy.create_team", lambda c: c.proxy.create_team(TeamNewBody(team_alias="owned"))), + ("proxy.delete_team", lambda c: c.proxy.delete_team("owned")), + ("proxy.delete_user", lambda c: c.proxy.delete_user("owned")), + ("proxy.spend_logs", lambda c: c.proxy.spend_logs(SpendLogsParams(api_key="owned"))), + ("proxy.probe", lambda c: c.proxy.probe("/user/info", params=NoBody())), +) + + +@pytest.mark.parametrize( + ("name", "operation"), MANAGEMENT_OPERATIONS, ids=tuple(name for name, _ in MANAGEMENT_OPERATIONS) +) +@pytest.mark.parametrize("kind", ("master", "direct_jwt", "virtual_key", "dashboard_session")) +def test_management_operations_send_the_selected_credential( + name: str, + operation: Callable[[ManagementClient], object], + kind: CredentialKind, +) -> None: + with caller_boundary(status=401) as (bootstrap, received), without_retries(): + client: Final = ( + bootstrap + if kind == "master" + else bootstrap.with_caller(Caller(credential=f"synthetic-{kind}", kind=kind, role="internal_user")) + ) + try: + operation(client) + except AssertionError: + pass + expected: Final = "Bearer bootstrap" if kind == "master" else f"Bearer synthetic-{kind}" + assert received.get_nowait() == expected, name + assert received.empty(), "an unauthorized request must not be retried" + + +class TestSplitCallerPropagation: + def test_control_and_data_replica_readers_keep_the_caller(self) -> None: + with caller_boundary() as (data, data_headers), caller_boundary() as (control, control_headers): + data_url: Final = next(iter(data.proxy.replicas)) + control_url: Final = next(iter(control.proxy.replicas)) + proxy: Final = build_proxy_client( + base_url=data_url, + control_plane_base_url=control_url, + replica_urls=(data_url,), + master_key="bootstrap", + ).with_caller(Caller(credential="tenant-token", kind="direct_jwt", role="team_member")) + proxy.key_info("owned") + proxy.read_body_back_everywhere( + "/key/info", KeyInfoResponse, settled=lambda info: info.info.key_alias == "owned" + ) + proxy.read_back_everywhere( + "/key/info", + params=NoBody(), + response_type=KeyInfoResponse, + converged=lambda result: isinstance(result, Success), + ) + assert control_headers.get_nowait() == "Bearer tenant-token" + assert control_headers.get_nowait() == "Bearer tenant-token" + assert data_headers.get_nowait() == "Bearer tenant-token" + assert control_headers.empty() and data_headers.empty() + + def test_successful_team_and_model_polling_uses_the_bound_caller(self) -> None: + with caller_boundary() as (bootstrap, received): + bound: Final = bootstrap.with_caller(Caller(credential="caller", kind="direct_jwt", role="proxy_admin")) + bound.create_team(TeamNewBody(team_alias="owned")) + bound.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic")) + assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer caller",) * 4 + assert received.empty() + + def test_expired_shaped_token_is_sent_once_without_renewal(self) -> None: + with caller_boundary(status=401) as (bootstrap, received): + bound: Final = bootstrap.with_caller( + Caller(credential="expired.payload.signature", kind="direct_jwt", role="internal_user") + ) + result: Final = bound.key_info_as("owned") + assert not isinstance(result, Success) + assert received.get_nowait() == "Bearer expired.payload.signature" + assert received.empty() + + +@pytest.mark.parametrize("operation", ("server", "toolset")) +def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: str) -> None: + bodies: Final[SimpleQueue[bytes]] = SimpleQueue() + with caller_boundary(status=401, bodies=bodies) as (bootstrap, _): + try: + if operation == "server": + bootstrap.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None)) + else: + bootstrap.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None)) + except AssertionError: + pass + expected: Final = ( + {"server_id": "owned", "alias": None} + if operation == "server" + else {"toolset_id": "owned", "description": None} + ) + assert json.loads(bodies.get_nowait()) == expected + assert bodies.empty() diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index e8caa801467..037db0c340f 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -7,11 +7,9 @@ client touches requests.* or builds raw dicts; they pass pydantic models here. from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Protocol -from pydantic import BaseModel - import e2e_http from e2e_http import ( URL, @@ -21,6 +19,7 @@ from e2e_http import ( Result, StreamingResponse, ) +from pydantic import BaseModel class Transport(Protocol): @@ -85,7 +84,7 @@ class Transport(Protocol): self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: ... - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ... + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: ... def upload[R: BaseModel]( self, @@ -113,7 +112,7 @@ class Transport(Protocol): @dataclass(frozen=True, slots=True) class HttpTransport: base_url: str - master_key: str + master_key: str = field(repr=False) request_timeout: float = 60.0 def _url(self, path: str) -> URL: @@ -245,10 +244,10 @@ class HttpTransport: timeout=self.request_timeout, ) - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), - headers=self.master, + headers=self.master if headers is None else headers, params=params, timeout=self.request_timeout, ) @@ -434,8 +433,8 @@ class SplitTransport: path, headers=headers, json=json, params=params, stream=stream ) - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - return self._route(path).probe(path, params=params) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: + return self._route(path).probe(path, params=params, headers=headers) def upload[R: BaseModel]( self, diff --git a/tests/e2e/ui/oidcSetup.ts b/tests/e2e/ui/oidcSetup.ts new file mode 100644 index 00000000000..943fa23cab3 --- /dev/null +++ b/tests/e2e/ui/oidcSetup.ts @@ -0,0 +1,30 @@ +import { chromium, expect } from "@playwright/test"; +import * as fs from "fs"; +import * as path from "path"; + +export default async function oidcSetup() { + const baseURL = process.env.E2E_OIDC_UI_URL; + const issuer = process.env.JWT_ISSUER; + const username = process.env.E2E_OIDC_USERNAME; + const password = process.env.E2E_OIDC_PASSWORD; + if (!baseURL || !issuer || !username || !password) { + throw new Error("The OIDC setup requires a running stack, issuer, and provisioned actor credentials"); + } + const artifactDir = process.env.E2E_UI_ARTIFACT_DIR || "."; + fs.mkdirSync(artifactDir, { recursive: true }); + const browser = await chromium.launch(); + try { + const page = await browser.newPage(); + await page.goto(`${baseURL.replace(/\/$/, "")}/sso/key/generate`); + await expect(page).toHaveURL(new RegExp(`^${issuer.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}/`)); + await page.getByLabel("Username or email").fill(username); + await page.getByLabel("Password", { exact: true }).fill(password); + await page.getByRole("button", { name: "Sign In", exact: true }).click(); + await page.waitForURL((url) => url.origin === new URL(baseURL).origin && url.pathname.startsWith("/ui")); + const statePath = path.join(artifactDir, "oidc.storageState.json"); + await page.context().storageState({ path: statePath }); + fs.chmodSync(statePath, 0o600); + } finally { + await browser.close(); + } +} diff --git a/tests/e2e/ui/playwright.oidc.config.ts b/tests/e2e/ui/playwright.oidc.config.ts new file mode 100644 index 00000000000..0fbe77e9bd2 --- /dev/null +++ b/tests/e2e/ui/playwright.oidc.config.ts @@ -0,0 +1,22 @@ +import { defineConfig, devices } from "@playwright/test"; +import * as path from "path"; + +const baseURL = process.env.E2E_OIDC_UI_URL; +if (!baseURL) throw new Error("E2E_OIDC_UI_URL must point to the running OIDC stack"); + +export default defineConfig({ + testDir: ".", + testMatch: "oidc/**/*.spec.ts", + retries: 0, + workers: 1, + outputDir: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc", "test-results"), + globalSetup: require.resolve("./oidcSetup"), + use: { + ...devices["Desktop Chrome"], + baseURL, + storageState: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc.storageState.json"), + trace: "off", + screenshot: "off", + video: "off", + }, +}); From c8bb54993e8b6db4eda84819b0015a1d9a85ca99 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:49:49 -0700 Subject: [PATCH 041/297] test: enforce isolated actors and stop OIDC process groups --- .../e2e/coverage_registry/management_cases.py | 2 +- tests/e2e/idp.py | 41 +++++++++++++++---- tests/e2e/management/jwt_actors.py | 2 +- tests/e2e/management/management_client.py | 17 ++++---- .../e2e/management/test_jwt_management_e2e.py | 20 ++++----- tests/e2e/test_idp.py | 25 ++++++++--- tests/e2e/test_proxy_client.py | 7 ++++ 7 files changed, 80 insertions(+), 34 deletions(-) diff --git a/tests/e2e/coverage_registry/management_cases.py b/tests/e2e/coverage_registry/management_cases.py index 15dc7d333c5..812dbfe5b8d 100644 --- a/tests/e2e/coverage_registry/management_cases.py +++ b/tests/e2e/coverage_registry/management_cases.py @@ -63,7 +63,7 @@ MANAGEMENT_CASES: Final = tuple( node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]", credential_kind="virtual_key", actor="proxy_admin", - profile="database_role", + profile="group_scoped", method="POST", path="/key/generate", operation_family="key_lifecycle", diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 12db91bbd88..2dc7c2ad71b 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -8,6 +8,7 @@ import secrets import signal import subprocess import sys +import time import warnings from collections.abc import Callable from contextlib import ExitStack @@ -426,6 +427,36 @@ def token_claims(token: str) -> TokenClaims: return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) +def _signal_process_group(process_id: int, signum: int) -> bool: + try: + os.killpg(process_id, signum) + except ProcessLookupError: + return False + return True + + +def _stop_process_group(child: subprocess.Popen[bytes]) -> None: + _signal_process_group(child.pid, signal.SIGTERM) + deadline: Final = time.monotonic() + 5 + while _process_group_exists(child.pid): + child.poll() + if time.monotonic() >= deadline: + _signal_process_group(child.pid, signal.SIGKILL) + break + time.sleep(0.05) + child.wait() + + +def _process_group_exists(process_id: int) -> bool: + try: + os.killpg(process_id, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + def run_oidc_profile(proxy_url: str, command: list[str]) -> int: idp: Final = keycloak_from_env().with_strict_cleanup() with ExitStack() as cleanup: @@ -441,17 +472,11 @@ def run_oidc_profile(proxy_url: str, command: list[str]) -> int: client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer) environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url} - with subprocess.Popen(command, env=environment) as child: + with subprocess.Popen(command, env=environment, start_new_session=True) as child: try: return child.wait() finally: - if child.poll() is None: - child.terminate() - try: - child.wait(timeout=5) - except subprocess.TimeoutExpired: - child.kill() - child.wait() + _stop_process_group(child) if __name__ == "__main__": diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py index 909d1652ada..2d23549fe71 100644 --- a/tests/e2e/management/jwt_actors.py +++ b/tests/e2e/management/jwt_actors.py @@ -84,7 +84,7 @@ class ActorFactory: ) ) ) - self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key)) + self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key, missing_ok=True)) return created def tenant(self) -> Tenant: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index a0243e868e2..8470d318db8 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -167,17 +167,18 @@ class ManagementClient: response_type=KeyInfoResponse, ) - def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> None: + def delete_key_strict(self, key: str, *, caller_key: str | None = None, missing_ok: bool = False) -> None: """Strict delete for the act phase of a test: a failed delete is a hard failure, unlike the warn-only ProxyClient.delete_key used at teardown.""" - _ = unwrap( - self.proxy.transport.post( - "/key/delete", - headers=self.proxy.management_headers(caller_key), - json=KeyDeleteBody(keys=[key]), - response_type=NoBody, - ) + result = self.proxy.transport.post( + "/key/delete", + headers=self.proxy.management_headers(caller_key), + json=KeyDeleteBody(keys=[key]), + response_type=NoBody, ) + if missing_ok and isinstance(result, UnknownApiError) and result.status_code == 404: + return + _ = unwrap(result) def delete_model_strict(self, model_id: str) -> None: """Strict delete for the act phase of a test: a failed delete is a hard diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py index 0d23f954158..5898073a4e6 100644 --- a/tests/e2e/management/test_jwt_management_e2e.py +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -102,31 +102,29 @@ class TestJwtManagement: @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key")) def test_admin_creates_reads_updates_clears_and_deletes_a_key( self, - client: ManagementClient, - idp: Keycloak, - jwt_identity: Identity, - resources: ResourceManager, actor_factory: ActorFactory, credential_kind: Literal["direct_jwt", "virtual_key"], ) -> None: - actor: Final = actor_factory.create("proxy_admin") + tenant: Final = actor_factory.tenant() + actor: Final = actor_factory.create("proxy_admin", tenants=(tenant,), profile="group_scoped") virtual_key: Final = ( actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None ) - admin: Final = ( - virtual_key if virtual_key is not None else idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + admin: Final = virtual_key if virtual_key is not None else actor.mint_caller(actor_factory.idp).credential + bound: Final = actor_factory.bootstrap.with_caller( + Caller(credential=admin, kind=credential_kind, role="proxy_admin") ) - bound: Final = client.with_caller(Caller(credential=admin, kind=credential_kind, role="proxy_admin")) + assert bound.user_info().user_id == actor.identity.user_id alias: Final = f"e2e-jwt-key-{unique_marker()}" created: Final = unwrap( bound.generate_key( - KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), + KeyGenerateBody(key_alias=alias, team_id=tenant.team_id, models=[CHEAP_OPENAI_MODEL]), ) ) - resources.defer(lambda: client.proxy.delete_key(created.key)) + actor_factory.resources.defer(lambda: actor_factory.bootstrap.delete_key_strict(created.key, missing_ok=True)) original: Final = unwrap(bound.key_info_as(created.key)).info - assert original.key_alias == alias and original.team_id == jwt_identity.group + assert original.key_alias == alias and original.team_id == tenant.team_id assert original.models == [CHEAP_OPENAI_MODEL] updated_alias: Final = f"{alias}-updated" diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py index 8a92dfa527c..cf2d4f3118a 100644 --- a/tests/e2e/test_idp.py +++ b/tests/e2e/test_idp.py @@ -6,6 +6,7 @@ from __future__ import annotations import os import signal +import socket import subprocess import sys import time @@ -148,16 +149,27 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No assert deletions.empty() -@pytest.mark.parametrize("exit_mode", ("normal", "parent", "group")) +@pytest.mark.parametrize( + ("exit_mode", "ignore_termination"), (("normal", False), ("parent", False), ("group", False), ("parent", True)) +) def test_oidc_launcher_removes_client_on_exit_and_termination( - tmp_path: Path, exit_mode: Literal["normal", "parent", "group"] + tmp_path: Path, exit_mode: Literal["normal", "parent", "group"], ignore_termination: bool ) -> None: ready: Final = tmp_path / "ready" + descendant_command: Final = ( + "import signal,socket,time; from pathlib import Path; " + + ("signal.signal(signal.SIGTERM, signal.SIG_IGN); " if ignore_termination else "") + + "listener=socket.socket(); listener.bind(('127.0.0.1',0)); listener.listen(); " + f"Path({str(ready)!r}).write_text(str(listener.getsockname()[1])); time.sleep(120)" + ) child_command: Final = ( - "import os,time; from pathlib import Path; " + "import os,subprocess,sys,time; from pathlib import Path; " 'assert os.environ["GENERIC_CLIENT_SECRET"]; ' 'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; ' - f"Path({str(ready)!r}).touch(); " + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") + f"subprocess.Popen([sys.executable, '-c', {descendant_command!r}]); " + f"ready=Path({str(ready)!r})\n" + "while not ready.exists(): time.sleep(0.05)\n" + + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") ) with _idp_server() as (idp, deletions): with subprocess.Popen( @@ -187,7 +199,10 @@ def test_oidc_launcher_removes_client_on_exit_and_termination( process.terminate() elif exit_mode == "group": os.killpg(process.pid, signal.SIGTERM) - assert process.wait(timeout=10) == (7 if exit_mode == "normal" else 143) + assert process.wait(timeout=15) == (7 if exit_mode == "normal" else 143) + with socket.socket() as connection: + connection.settimeout(1) + assert connection.connect_ex(("127.0.0.1", int(ready.read_text()))) != 0 finally: if process.poll() is None: os.killpg(process.pid, signal.SIGKILL) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 879bd88980c..0c4aed5bd65 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -121,6 +121,13 @@ def caller_boundary( class TestBoundManagementCaller: + def test_strict_key_cleanup_accepts_missing_only_when_requested(self) -> None: + with caller_boundary(delete_status=404) as (bootstrap, received), without_retries(): + with pytest.raises(AssertionError): + bootstrap.delete_key_strict("owned") + bootstrap.delete_key_strict("owned", missing_ok=True) + assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap") + def test_actor_key_cleanup_reports_failure_and_continues(self) -> None: with caller_boundary(delete_status=500) as (bootstrap, received), without_retries(): resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True) From cba843cc167bbadf5d18bee0ea07a443a4838ba1 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 12 Sep 2026 12:03:04 -0700 Subject: [PATCH 042/297] feat(proxy): predict prompt-cache costs across deployments --- .../llms/anthropic/prompt_cache_prediction.py | 388 ++++++++++ litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 5 +- litellm/proxy/auth/auth_utils.py | 35 +- litellm/proxy/auth/user_api_key_auth.py | 15 +- .../common_utils/prompt_cache_pricing.py | 91 +++ litellm/proxy/hooks/__init__.py | 2 + .../hooks/parallel_request_limiter_v3.py | 204 ++--- .../proxy/hooks/prompt_cache_prediction.py | 142 ++++ .../cost_tracking_settings.py | 2 + .../prompt_cache_prediction.py | 278 +++++++ .../streaming_handler.py | 18 + .../prompt_cache_prediction.py | 67 ++ .../test_anthropic_prompt_cache_prediction.py | 209 ++++++ .../proxy/auth/test_auth_utils.py | 226 ++++++ .../common_utils/test_prompt_cache_pricing.py | 105 +++ .../hooks/test_parallel_request_limiter_v3.py | 245 ++++++ .../proxy/hooks/test_prompt_cache_observer.py | 300 ++++++++ .../test_prompt_cache_prediction.py | 698 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 164 ++++ 20 files changed, 3106 insertions(+), 89 deletions(-) create mode 100644 litellm/llms/anthropic/prompt_cache_prediction.py create mode 100644 litellm/proxy/common_utils/prompt_cache_pricing.py create mode 100644 litellm/proxy/hooks/prompt_cache_prediction.py create mode 100644 litellm/proxy/management_endpoints/prompt_cache_prediction.py create mode 100644 litellm/types/management_endpoints/prompt_cache_prediction.py create mode 100644 tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py create mode 100644 tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py create mode 100644 tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_prompt_cache_prediction.py diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py new file mode 100644 index 00000000000..e69a02bd93a --- /dev/null +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -0,0 +1,388 @@ +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate +from types import MappingProxyType +from typing import Annotated, Final, Literal, Protocol, TypeAlias + +import httpx +from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError + +import litellm +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION +from litellm.types.router import LiteLLM_Params +from litellm.types.utils import ModelResponse + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_HEADERS: Final = TypeAdapter(dict[str, str]) +_counter: Final = AnthropicCountTokensHandler() + + +_NATIVE_HEADERS: Final = frozenset( + ( + "host", + "accept", + "accept-encoding", + "connection", + "user-agent", + "content-length", + "content-type", + "x-api-key", + "anthropic-version", + ) +) + +_DEPLOYMENT_OPTIONS: Final = frozenset( + { + "model", + "api_key", + "api_base", + "custom_llm_provider", + "rpm", + "tpm", + "timeout", + "stream_timeout", + "max_retries", + "num_retries", + "max_parallel_requests", + "input_cost_per_token", + "output_cost_per_token", + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + } +) + + +class _StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + +class _CacheControl(_StrictModel): + type: Literal["ephemeral"] + ttl: Literal["5m", "1h"] = "5m" + + +class _Text(_StrictModel): + type: Literal["text"] + text: str = Field(min_length=1, pattern=r"\S") + cache_control: _CacheControl | None = None + + +class _ToolUse(_StrictModel): + type: Literal["tool_use"] + id: str = Field(min_length=1) + name: str = Field(min_length=1) + input: Mapping[str, JsonValue] + cache_control: _CacheControl | None = None + + +class _ResultText(_StrictModel): + type: Literal["text"] + text: str + + +class _ToolResult(_StrictModel): + type: Literal["tool_result"] + tool_use_id: str = Field(min_length=1) + content: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] + is_error: bool | None = None + cache_control: _CacheControl | None = None + + +_Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminator="type")] + + +class _Message(_StrictModel): + role: Literal["user", "assistant"] + content: str | Annotated[tuple[_Block, ...], Field(strict=False)] + + def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]: + return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content) + + +class _Tool(_StrictModel): + name: str = Field(min_length=1) + description: str | None = None + input_schema: Mapping[str, JsonValue] + type: Literal["custom"] | None = None + + +class _Request(_StrictModel): + messages: tuple[_Message, ...] = Field(min_length=1, strict=False) + system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None + tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None + model: str | None = None + max_tokens: int | None = None + stream: bool | None = None + temperature: float | int | None = None + top_p: float | int | None = None + top_k: int | None = None + stop_sequences: Annotated[tuple[str, ...], Field(strict=False)] | None = None + metadata: Mapping[str, JsonValue] | None = None + + +@dataclass(frozen=True, slots=True) +class PromptPrefix: + prefix_body: Mapping[str, JsonValue] + fingerprint: str + fingerprints: tuple[str, ...] + ttl_seconds: int + + +def _digest(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode() + ).hexdigest() + + +def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str: + return _digest((previous, boundary)) + + +def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None: + try: + request: Final = _Request.model_validate(body) + blocks: Final = tuple(message.blocks() for message in request.messages) + except ValidationError: + return None + markers: Final = tuple( + (message_index, block_index, block.cache_control) + for message_index, message_blocks in enumerate(blocks) + for block_index, block in enumerate(message_blocks) + if block.cache_control is not None + ) + if len(markers) != 1: + return None + message_end, block_end, marker = markers[0] + normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True)) + context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized}) + boundaries: Final = tuple( + ( + message_index, + request.messages[message_index].role, + _JSON_OBJECT.validate_python( + block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True) + ), + ) + for message_index, message_blocks in enumerate(blocks[: message_end + 1]) + for block_index, block in enumerate(message_blocks) + if message_index < message_end or block_index <= block_end + ) + hashes: Final = tuple( + accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl))) + )[1:] + prefix_messages: Final = tuple( + _Message( + role=request.messages[message_index].role, + content=tuple( + block + for block_index, block in enumerate(message_blocks) + if message_index < message_end or block_index <= block_end + ), + ) + for message_index, message_blocks in enumerate(blocks[: message_end + 1]) + ) + return PromptPrefix( + prefix_body=MappingProxyType( + _JSON_OBJECT.validate_python( + _Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump( + mode="json", exclude_none=True + ) + ) + ), + fingerprint=hashes[-1], + fingerprints=tuple(reversed(hashes[-20:])), + ttl_seconds=3600 if marker.ttl == "1h" else 300, + ) + + +def cache_scope( + caller_key_hash: str, + deployment_id: str, + provider_key: str, + model: str, + anthropic_version: str = DEFAULT_ANTHROPIC_API_VERSION, +) -> str: + return _digest((caller_key_hash, deployment_id, provider_key, model, anthropic_version)) + + +class _TTLUsage(BaseModel): + model_config = ConfigDict(strict=True) + ephemeral_5m_input_tokens: int = Field(default=0, ge=0) + ephemeral_1h_input_tokens: int = Field(default=0, ge=0) + + +class _CacheUsage(BaseModel): + model_config = ConfigDict(strict=True) + cached_tokens: int = Field(default=0, ge=0) + cache_creation_tokens: int = Field(default=0, ge=0) + cache_creation_token_details: _TTLUsage | None = None + + +class _Usage(BaseModel): + model_config = ConfigDict(strict=True) + prompt_tokens: int = Field(ge=0) + prompt_tokens_details: _CacheUsage + + +class _Choice(BaseModel): + finish_reason: str = Field(min_length=1) + + +class _Response(BaseModel): + model_config = ConfigDict(strict=True) + model: str + usage: _Usage + choices: tuple[_Choice, ...] = Field(min_length=1, strict=False) + + +class _CountBody(BaseModel): + messages: Sequence[Mapping[str, JsonValue]] + tools: Sequence[Mapping[str, JsonValue]] | None = None + system: str | Sequence[Mapping[str, JsonValue]] | None = None + + +class _CountResult(BaseModel): + input_tokens: Annotated[StrictInt, Field(ge=0)] + + +class TokenCounter(Protocol): + async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: ... + + +def _count_objects( + values: Sequence[Mapping[str, JsonValue]], +) -> list[dict[str, JsonValue]]: # mutable-ok: the existing provider count API requires JSON lists/dicts + return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary + + +async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + native: Final = _CountBody.model_validate(body) + try: + result: Final = _CountResult.model_validate( + await _counter.handle_count_tokens_request( + model=model, + messages=_count_objects(native.messages), + tools=_count_objects(native.tools) if native.tools is not None else None, + system=native.system, + api_key=api_key, + timeout=15.0, + ) + ) + except Exception: # noqa: BLE001 # provider/count validation failures are unavailable estimates, not zero tokens + return None + return result.input_tokens + + +@dataclass(frozen=True, slots=True) +class NativePredictionTarget: + model: str + api_key: str + + +@dataclass(frozen=True, slots=True) +class UnsupportedPredictionTarget: + reason: Literal[ + "unsupported_deployment_configuration", + "unsupported_provider_endpoint", + "unsupported_provider", + "unsupported_provider_credentials", + ] + + +def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget: + configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True)) + if configured_options - _DEPLOYMENT_OPTIONS: + return UnsupportedPredictionTarget("unsupported_deployment_configuration") + api_base: Final = AnthropicModelInfo.get_api_base(params.api_base) + if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"): + return UnsupportedPredictionTarget("unsupported_provider_endpoint") + try: + model, provider, _, _ = litellm.get_llm_provider( + model=params.model, custom_llm_provider=params.custom_llm_provider + ) + except Exception: # noqa: BLE001 # the shared provider resolver raises for unknown deployments + return UnsupportedPredictionTarget("unsupported_provider") + if provider != "anthropic": + return UnsupportedPredictionTarget("unsupported_provider") + api_key: Final = AnthropicModelInfo.get_api_key(params.api_key) + if api_key is None or not _supported_provider_key(api_key): + return UnsupportedPredictionTarget("unsupported_provider_credentials") + return NativePredictionTarget(model=model, api_key=api_key) + + +def _supported_provider_key(api_key: str) -> bool: + return bool(api_key) and not is_anthropic_oauth_key(api_key) + + +def supported_prediction_headers(headers: Mapping[str, str]) -> bool: + return all( + name.lower() != "anthropic-beta" + and (name.lower() != "anthropic-version" or value == DEFAULT_ANTHROPIC_API_VERSION) + for name, value in headers.items() + ) + + +@dataclass(frozen=True, slots=True) +class ObservedCachePrefix: + prefix: PromptPrefix + scope: str + cached_tokens: int + cache_creation_tokens: int + + +def parse_observed_cache( + wire: httpx.Request, response_obj: ModelResponse, caller_key_hash: str, deployment_id: str +) -> ObservedCachePrefix | None: + try: + response: Final = _Response.model_validate(response_obj, from_attributes=True) + body: Final = _JSON_OBJECT.validate_json(wire.content) + headers: Final = _HEADERS.validate_python(wire.headers) + except (ValidationError, RuntimeError, httpx.RequestNotRead): + return None + if ( + wire.url.scheme != "https" + or wire.url.host != "api.anthropic.com" + or wire.url.path != "/v1/messages" + or wire.url.query + or wire.url.port not in (None, 443) + ): + return None + if ( + frozenset(headers) - _NATIVE_HEADERS + or not supported_prediction_headers(headers) + or headers.get("anthropic-version") != DEFAULT_ANTHROPIC_API_VERSION + ): + return None + provider_key: Final = headers.get("x-api-key", "") + model: Final = body.get("model") + if not _supported_provider_key(provider_key) or not isinstance(model, str) or model != response.model: + return None + prefix: Final = parse_prompt(body) + if prefix is None: + return None + usage: Final = response.usage.prompt_tokens_details + cache_tokens: Final = usage.cached_tokens + usage.cache_creation_tokens + if cache_tokens <= 0 or cache_tokens > response.usage.prompt_tokens: + return None + split: Final = usage.cache_creation_token_details + if usage.cache_creation_tokens and split is None: + return None + if split is not None and ( + split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens != usage.cache_creation_tokens + or (prefix.ttl_seconds == 300 and split.ephemeral_1h_input_tokens > 0) + or (prefix.ttl_seconds == 3600 and split.ephemeral_5m_input_tokens > 0) + ): + return None + return ObservedCachePrefix( + prefix=prefix, + scope=cache_scope(caller_key_hash, deployment_id, provider_key, model), + cached_tokens=cache_tokens, + cache_creation_tokens=usage.cache_creation_tokens, + ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ae6c042ab3a..de82d8ec3c8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -887,6 +887,7 @@ class LiteLLMRoutes(enum.Enum): "/auto_router/validate_complexity_router_config", # Per-session auto-router read - the endpoint scopes the row to the caller's own key hash "/auto_router/session", + "/cost/predict-cache", # Agent registry - reads are role-scoped and writes are proxy-admin-gated # inside agent_endpoints/endpoints.py *agent_management_routes, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9c175242a9a..3495bf2ae98 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -878,6 +878,7 @@ async def common_checks( request_query_params=_safe_get_request_query_params(request=request), llm_router=llm_router, request=request, + team_id=valid_token.team_id if valid_token is not None else None, ) skip_all_budget_checks: Final = skip_budget_checks or ( @@ -4356,7 +4357,7 @@ async def stamp_matched_model_access_groups( async def can_key_call_model( model: str | list[str], - llm_model_list: list | None, + llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: litellm.Router | None, ) -> Literal[True]: @@ -4403,7 +4404,7 @@ async def can_key_call_model( async def can_key_call_resolved_model( model: str, - llm_model_list: list | None, + llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: litellm.Router | None, ) -> None: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index be65c3b39ec..dc304a156cf 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -33,7 +33,7 @@ from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_me from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) -from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS +from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment from litellm.types.utils import CustomPricingLiteLLMParams @@ -1736,7 +1736,7 @@ def _append_model_candidates(candidates: list[str], value: Any) -> None: candidates.extend(model for model in model_names if model) -def _dedupe_model_candidates(candidates: list[str]) -> list[str]: +def _dedupe_model_candidates(candidates: Collection[str]) -> list[str]: deduped: Final[list[str]] = [] for model in candidates: if model not in deduped: @@ -1845,13 +1845,42 @@ def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | Non return model_id +def get_cache_prediction_deployments( + *, current_deployment_id: str, candidate_deployment_id: str, llm_router: Router, team_id: str | None +) -> tuple[Deployment, Deployment] | None: + current: Final = llm_router.get_deployment(current_deployment_id) + candidate: Final = llm_router.get_deployment(candidate_deployment_id) + if current is None or candidate is None: + return None + if any(deployment.model_info.team_id not in (None, team_id) for deployment in (current, candidate)): + return None + return current, candidate + + +def _cache_prediction_model_candidates( + request_data: Mapping[str, object], llm_router: Router | None, team_id: str | None +) -> tuple[str, ...]: + current_id: Final = request_data.get("current_deployment_id") + candidate_id: Final = request_data.get("candidate_deployment_id") + if llm_router is None or not isinstance(current_id, str) or not isinstance(candidate_id, str): + return () + deployments: Final = get_cache_prediction_deployments( + current_deployment_id=current_id, candidate_deployment_id=candidate_id, llm_router=llm_router, team_id=team_id + ) + return tuple(deployment.model_name for deployment in deployments) if deployments is not None else () + + def _extract_model_candidates_from_request( request_data: dict, route: str, request_headers: Mapping[str, object] | None = None, request_query_params: Mapping[str, object] | None = None, llm_router: Router | None = None, + team_id: str | None = None, ) -> list[str]: + if route == "/cost/predict-cache": + prediction_models: Final = _cache_prediction_model_candidates(request_data, llm_router, team_id) # pyright: ignore[reportUnknownArgumentType] # the typed reader validates each deployment ID from this legacy payload + return _dedupe_model_candidates(prediction_models) candidates: Final[list[str]] = [] uses_model_routing_sources: Final = _route_uses_model_routing_sources(route=route) uses_header_or_query_model_sources: Final = _route_matches_any_marker( @@ -1945,6 +1974,7 @@ def get_model_from_request( request_query_params: Mapping[str, object] | None = None, llm_router: Router | None = None, request: Request | None = None, + team_id: str | None = None, ) -> str | list[str] | None: """Resolve the model(s) a request targets, for model-access and budget checks. @@ -1967,6 +1997,7 @@ def get_model_from_request( request_headers=request_headers, request_query_params=request_query_params, llm_router=llm_router, + team_id=team_id, ) model = _format_model_candidates(candidates) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9828311112e..930e3cca703 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -182,6 +182,7 @@ def _get_model_from_request_context( route: str, request: Request | None, llm_router: Any | None = None, + team_id: str | None = None, ) -> str | list[str] | None: return get_model_from_request( request_data=request_data, @@ -190,6 +191,7 @@ def _get_model_from_request_context( request_query_params=_safe_get_request_query_params(request=request), llm_router=llm_router, request=request, + team_id=team_id, ) @@ -208,7 +210,7 @@ async def _normalize_claude_model( return if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True: return - requested: Final = _get_model_from_request_context(request_data, route, request, llm_router) + requested: Final = _get_model_from_request_context(request_data, route, request, llm_router, valid_token.team_id) if not isinstance(requested, str) or requested != request_data.get("model"): return if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): @@ -1592,6 +1594,7 @@ async def _user_api_key_auth_builder( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) skip_budget_checks = False if model is not None and llm_router is not None: @@ -1632,6 +1635,7 @@ async def _user_api_key_auth_builder( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) ), ) @@ -2022,6 +2026,7 @@ async def _user_api_key_auth_builder( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) skip_budget_checks = False if model is not None and llm_router is not None: @@ -2140,6 +2145,7 @@ async def _user_api_key_auth_builder( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) current_models = _get_model_names_for_budget_checks(model=current_model) @@ -2170,6 +2176,7 @@ async def _user_api_key_auth_builder( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) current_models = _get_model_names_for_budget_checks(model=current_model) @@ -2665,6 +2672,7 @@ async def _run_centralized_common_checks( route=route, request=request, llm_router=llm_router, + team_id=user_api_key_auth_obj.team_id, ) # Pin the metadata variable name (litellm_metadata vs metadata) before @@ -2781,12 +2789,14 @@ def _should_skip_budget_checks( route: str, request: Request | None, llm_router: Any | None, + team_id: str | None = None, ) -> bool: model: Final = _get_model_from_request_context( request_data=request_data, route=route, request=request, llm_router=llm_router, + team_id=team_id, ) if model is not None and llm_router is not None: return _is_model_cost_zero(model=model, llm_router=llm_router) @@ -3232,6 +3242,7 @@ async def _enforce_key_and_fallback_model_access( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) if model is not None: @@ -3339,6 +3350,7 @@ async def _run_post_custom_auth_checks( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) current_models = _get_model_names_for_budget_checks(model=current_model) @@ -3380,6 +3392,7 @@ async def _run_post_custom_auth_checks( route=route, request=request, llm_router=llm_router, + team_id=valid_token.team_id, ) current_models = _get_model_names_for_budget_checks(model=current_model) diff --git a/litellm/proxy/common_utils/prompt_cache_pricing.py b/litellm/proxy/common_utils/prompt_cache_pricing.py new file mode 100644 index 00000000000..ff070853b46 --- /dev/null +++ b/litellm/proxy/common_utils/prompt_cache_pricing.py @@ -0,0 +1,91 @@ +from collections.abc import Mapping +from math import isfinite +from typing import Final + +from pydantic import TypeAdapter + +import litellm +from litellm.cost_calculator import ( + _select_model_name_for_cost_calc, # pyright: ignore[reportPrivateUsage] # shares completion_cost's deployment tariff selection + completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped +) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets +from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage + +_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object]) + + +def _valid_price(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0 + + +def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets) -> bool: + required: Final = ( + ("input_cost_per_token", True), + ("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0), + ("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0), + ("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0), + ) + if any(needed and not _valid_price(prices.get(key)) for key, needed in required): + return False + return all( + _valid_price(value) + for key, value in prices.items() + if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required) + ) + + +def price_cache_tokens(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> float | None: + try: + selected_model: Final = _select_model_name_for_cost_calc( + model=model, + completion_response=None, + custom_pricing=True, + custom_llm_provider="anthropic", + router_model_id=deployment_id, + ) + if selected_model is None: + return None + model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic") + registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary + price_entry: Final = registry.get(model_info["key"]) + if price_entry is None: + return None + prices: Final = _PRICE_ENTRY.validate_python(price_entry) + if not _has_required_prices(prices, tokens): + return None + usage: Final = Usage( + prompt_tokens=tokens.total_tokens, + completion_tokens=0, + total_tokens=tokens.total_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=tokens.cache_read_input_tokens, + cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens, + ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens, + ), + ), + ) + logging_obj: Final = Logging( + model=model, + messages=[], # mutable-ok: Logging requires a list + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="prompt-cache-prediction", + function_id="prompt-cache-prediction", + ) + completion_cost( + completion_response=ModelResponse(model=model, usage=usage), + model=model, + custom_llm_provider="anthropic", + custom_pricing=True, + router_model_id=deployment_id, + litellm_logging_obj=logging_obj, + ) + cost: Final = logging_obj.cost_breakdown.get("input_cost") if logging_obj.cost_breakdown is not None else None + return cost if cost is not None and _valid_price(cost) else None + except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models + return None diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 8714dd5f3d2..f3542098f95 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -9,6 +9,7 @@ from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler from .max_iterations_limiter import _PROXY_MaxIterationsHandler from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 +from .prompt_cache_prediction import PromptCacheObserver from .responses_id_security import ResponsesIDSecurity from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler @@ -25,6 +26,7 @@ PROXY_HOOKS: Final = { "max_iterations_limiter": _PROXY_MaxIterationsHandler, "max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler, "sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler, + "prompt_cache_prediction": PromptCacheObserver, } ## FEATURE FLAG HOOKS ## diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c398abff099..a34dc99e472 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -9,10 +9,12 @@ import binascii import logging import os import uuid -from collections.abc import Awaitable, Callable, Mapping, Sequence, Set +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set +from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime +from types import MappingProxyType from typing import ( TYPE_CHECKING, Any, @@ -23,6 +25,7 @@ from typing import ( TypedDict, ) +from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly from litellm import DualCache @@ -84,6 +87,9 @@ else: InternalUsageCache = Any +_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object]) + + BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) @@ -2673,12 +2679,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns list of descriptors for API key, user, team, team member, end user, model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import ( - get_team_model_rpm_limit, - get_team_model_tpm_limit, - ) - - descriptors: Final = [] + descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: existing descriptor helpers append in place # API Key rate limits if user_api_key_dict.api_key and ( @@ -2803,34 +2804,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors=descriptors, ) - if ( - get_team_model_rpm_limit(user_api_key_dict) is not None - or get_team_model_tpm_limit(user_api_key_dict) is not None - ): - _tpm_limit_for_team_model: Final = get_team_model_tpm_limit(user_api_key_dict) or {} - _rpm_limit_for_team_model: Final = get_team_model_rpm_limit(user_api_key_dict) or {} - should_check_rate_limit = False - if requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model: - should_check_rate_limit = True - - if should_check_rate_limit: - model_specific_tpm_limit = None - model_specific_rpm_limit = None - if requested_model in _tpm_limit_for_team_model: - model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model] - if requested_model in _rpm_limit_for_team_model: - model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model] - descriptors.append( - RateLimitDescriptor( - key="model_per_team", - value=f"{user_api_key_dict.team_id}:{requested_model}", - rate_limit={ - "requests_per_unit": model_specific_rpm_limit, - "tokens_per_unit": model_specific_tpm_limit, - "window_size": self.window_size, - }, - ) - ) + self._add_team_model_rate_limit_descriptor_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model if isinstance(requested_model, str) else None, + descriptors=descriptors, + ) # Agent-level and session-level rate limits resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data) @@ -3416,6 +3394,108 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model, ) + async def _build_request_rate_limit_descriptors( + self, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str | None, + ) -> list[RateLimitDescriptor]: # mutable-ok: the shared generation reservation helpers require a list + metadata: Final = _REQUEST_RATE_LIMIT_DATA.validate_python( + user_api_key_dict.metadata or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # validates the legacy auth metadata boundary + ) + rpm_value: Final = metadata.get("rpm_limit_type") + tpm_value: Final = metadata.get("tpm_limit_type") + rpm_limit_type: Final = rpm_value if isinstance(rpm_value, str) else None + tpm_limit_type: Final = tpm_value if isinstance(tpm_value, str) else None + model_value: Final = data.get("model") + requested_model: Final = model_value if isinstance(model_value, str) else None + model_has_failures: Final = ( + await self._check_model_has_recent_failures( + model=requested_model, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + if requested_model and self._is_dynamic_rate_limiting_enabled(rpm_limit_type, tpm_limit_type) + else False + ) + descriptors: Final = self._create_rate_limit_descriptors( # pyright: ignore[reportUnknownMemberType] # legacy helper reads a dictionary with validated keys + user_api_key_dict=user_api_key_dict, + data=dict(data), # mutable-ok: legacy descriptor helpers accept a request dictionary + rpm_limit_type=rpm_limit_type, + tpm_limit_type=tpm_limit_type, + model_has_failures=model_has_failures, + call_type=call_type, + ) + self._add_project_model_rate_limit_descriptor_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + descriptors=descriptors, + ) + self.add_project_io_token_rate_limit_descriptors_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + descriptors=descriptors, + ) + return [ # mutable-ok: the shared generation reservation helpers require a list + *descriptors, + *self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model), + ] + + async def _release_request_capacity_when_admitted( + self, + admission: asyncio.Task[RateLimitResponse], + acquisition: ParallelSlotAcquisition, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + response: Final = await admission + if response["overall_code"] == "OK": + await self._release_parallel_request_slots(acquisition, user_api_key_dict.parent_otel_span) + + @asynccontextmanager + async def request_capacity( + self, + user_api_key_dict: UserAPIKeyAuth, + model: str, + *, + request_data: Mapping[str, object] | None = None, + ) -> AsyncGenerator[None, None]: + """Charge one non-generation provider request to RPM and hold its concurrency slot.""" + data: Final = MappingProxyType({**(request_data or MappingProxyType({})), "model": model}) + descriptors: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None) + acquisition: Final = ParallelSlotAcquisition( + slot_id=uuid.uuid4().hex, + counter_keys=[ # mutable-ok: the shared slot-release contract requires a list + self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") + for d in descriptors + if d["rate_limit"] is not None and d["rate_limit"].get("max_parallel_requests") is not None + ], + ) + admission: Final = asyncio.create_task( + self.should_rate_limit( + descriptors=descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + skip_tpm_check=True, + parallel_slot_id=acquisition["slot_id"], + ) + ) + try: + response: Final = await asyncio.shield(admission) + if response["overall_code"] == "OVER_LIMIT": + self._handle_rate_limit_error(response, descriptors, model) + yield + finally: + cleanup: Final = asyncio.create_task( + self._release_request_capacity_when_admitted(admission, acquisition, user_api_key_dict) + ) + cancellation: asyncio.CancelledError | None = None # rebind-ok: retain cancellation until cleanup finishes + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + cancellation = exc # rebind-ok: retain the latest cancellation without interrupting slot release + cleanup.result() + if cancellation is not None: + raise cancellation + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -3444,59 +3524,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): call_type=call_type, ) - # Get rate limit types from metadata - metadata: Final = user_api_key_dict.metadata or {} - rpm_limit_type: Final = metadata.get("rpm_limit_type") - tpm_limit_type: Final = metadata.get("tpm_limit_type") - - # For dynamic mode, check if the model has recent failures - model_has_failures = False - requested_model: Final = data.get("model", None) - - if ( - self._is_dynamic_rate_limiting_enabled( - rpm_limit_type=rpm_limit_type, - tpm_limit_type=tpm_limit_type, - ) - and requested_model - ): - model_has_failures = await self._check_model_has_recent_failures( - model=requested_model, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - - # Create rate limit descriptors - descriptors: Final = self._create_rate_limit_descriptors( + request_data: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data) + model_value: Final = request_data.get("model") + requested_model: Final = model_value if isinstance(model_value, str) else None + descriptors: Final = await self._build_request_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, - data=data, - rpm_limit_type=rpm_limit_type, - tpm_limit_type=tpm_limit_type, - model_has_failures=model_has_failures, + data=request_data, call_type=call_type, ) - # Add team model rate limits from team_metadata - self._add_team_model_rate_limit_descriptor_from_metadata( - user_api_key_dict=user_api_key_dict, - requested_model=requested_model, - descriptors=descriptors, - ) - - # Project Level Rate Limits - self._add_project_model_rate_limit_descriptor_from_metadata( - user_api_key_dict=user_api_key_dict, - requested_model=requested_model, - descriptors=descriptors, - ) - self.add_project_io_token_rate_limit_descriptors_from_metadata( - user_api_key_dict=user_api_key_dict, - requested_model=requested_model, - descriptors=descriptors, - ) - - # Org Level Rate Limits - descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model)) - # Only check rate limits if we have descriptors with actual limits if descriptors: # First pass: RPM and max_parallel_requests sliding-window check. diff --git a/litellm/proxy/hooks/prompt_cache_prediction.py b/litellm/proxy/hooks/prompt_cache_prediction.py new file mode 100644 index 00000000000..65c456c5666 --- /dev/null +++ b/litellm/proxy/hooks/prompt_cache_prediction.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable, Mapping +from datetime import datetime +from typing import TYPE_CHECKING, Final, Literal + +import httpx +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, parse_observed_cache +from litellm.types.utils import ModelResponse + +if TYPE_CHECKING: + from litellm.proxy.utils import InternalUsageCache + +_RETENTION_SECONDS: Final = 86_400 + + +class CacheObservation(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$") + cached_tokens: int = Field(gt=0) + observed_at: float = Field(ge=0, allow_inf_nan=False) + expires_at: float = Field(ge=0, allow_inf_nan=False) + + +_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None) + + +def _cache_key(scope: str, fingerprint: str) -> str: + return f"prompt-cache-observation:{scope}:{fingerprint}" + + +async def lookup( + cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None +) -> CacheObservation | None: + checked_at: Final = time.time() if now is None else now + exact: Final = await _read_exact(cache, scope, prefix.fingerprint) + if exact is not None and exact.expires_at > checked_at: + return exact + older: Final = await asyncio.gather( + *(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:]) + ) + observations: Final = tuple(observation for observation in (exact, *older) if observation is not None) + return next( + (observation for observation in observations if observation.expires_at > checked_at), + next(iter(observations), None), + ) + + +async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None: + try: + value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # validate the legacy cache's untyped result at the I/O boundary + if value is None: + return None + observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value + except ValidationError: + return None + return observation if observation.fingerprint == fingerprint else None + + +class _Metadata(BaseModel): + model_config = ConfigDict(strict=True) + user_api_key_hash: str = Field(min_length=1) + + +class _Logged(BaseModel): + model_config = ConfigDict(strict=True) + status: Literal["success"] + model_id: str = Field(min_length=1) + metadata: _Metadata + + +class _Event(BaseModel): + model_config = ConfigDict(strict=True, arbitrary_types_allowed=True) + call_type: Literal["anthropic_messages"] + custom_llm_provider: Literal["anthropic"] + cache_hit: bool | None = None + httpx_response: httpx.Response + first_api_call_start_time: datetime + standard_logging_object: _Logged + stream: bool = False + prompt_cache_response_complete: bool = False + + +class PromptCacheObserver(CustomLogger): + def __init__(self, internal_usage_cache: InternalUsageCache, clock: Callable[[], float] = time.time) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # base callback constructor accepts untyped kwargs + self.cache = internal_usage_cache.dual_cache + self.clock = clock + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if not isinstance(response_obj, ModelResponse): + return + try: + event: Final = _Event.model_validate(kwargs) + wire: Final = event.httpx_response.request + except (ValidationError, RuntimeError, httpx.RequestNotRead): + return + if ( + event.cache_hit + or event.httpx_response.status_code != 200 + or (event.stream and not event.prompt_cache_response_complete) + ): + return + observed: Final = parse_observed_cache( + wire, + response_obj, + event.standard_logging_object.metadata.user_api_key_hash, + event.standard_logging_object.model_id, + ) + if observed is None: + return + prefix: Final = observed.prefix + scope: Final = observed.scope + cache_tokens: Final = observed.cached_tokens + now: Final = self.clock() + started: Final = event.first_api_call_start_time.timestamp() + if started > now: + return + if observed.cache_creation_tokens == 0: + previous: Final = await _read_exact(self.cache, scope, prefix.fingerprint) + if previous is None or previous.fingerprint != prefix.fingerprint or previous.cached_tokens != cache_tokens: + return + observation: Final = CacheObservation( + fingerprint=prefix.fingerprint, + cached_tokens=cache_tokens, + observed_at=now, + expires_at=started + prefix.ttl_seconds, + ) + key: Final = _cache_key(scope, prefix.fingerprint) + payload: Final = observation.model_dump_json() + await self.cache.async_set_cache(key, payload, ttl=_RETENTION_SECONDS) # pyright: ignore[reportUnknownMemberType] # legacy cache accepts a serialized validated observation + if self.cache.redis_cache is not None: + await self.cache.async_set_cache(key, payload, local_only=True, ttl=1) # pyright: ignore[reportUnknownMemberType] # keep the local copy short-lived while Redis retains stale evidence diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index dc0da63555f..cb376f286ec 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -28,6 +28,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.prompt_cache_prediction import router as prompt_cache_prediction_router from litellm.types.utils import ( CostBreakdown, CostPerToken, @@ -39,6 +40,7 @@ from litellm.types.utils import ( ) router: Final = APIRouter() +router.include_router(prompt_cache_prediction_router) @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/management_endpoints/prompt_cache_prediction.py b/litellm/proxy/management_endpoints/prompt_cache_prediction.py new file mode 100644 index 00000000000..56e844214d6 --- /dev/null +++ b/litellm/proxy/management_endpoints/prompt_cache_prediction.py @@ -0,0 +1,278 @@ +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Annotated, Final + +from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import BaseModel, JsonValue, TypeAdapter + +import litellm +from litellm._internal_context import current_billing_time, pinned_billing_time +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.anthropic.prompt_cache_prediction import ( + PromptPrefix, + TokenCounter, + UnsupportedPredictionTarget, + cache_scope, + count_prompt_tokens, + parse_prompt, + resolve_prediction_target, + supported_prediction_headers, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model +from litellm.proxy.auth.auth_utils import get_cache_prediction_deployments +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.http_parsing_utils import ( + _read_request_body, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # canonical parsed-body owner; validate its legacy result at the endpoint boundary +) +from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # use the configured proxy limiter's shared capacity owner +) +from litellm.proxy.hooks.prompt_cache_prediction import lookup +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.types.management_endpoints.prompt_cache_prediction import ( + CacheCostScenario, + CacheEvidence, + CachePredictionArm, + CachePredictionRequest, + CachePredictionResponse, + CacheTokenBuckets, +) +from litellm.types.router import Deployment +from litellm.utils import get_prompt_cache_min_tokens + +router: Final = APIRouter() +_REQUEST_DATA: Final = TypeAdapter(Mapping[str, object]) + + +class _CallerSettings(BaseModel): + config: Mapping[str, object] | None = None + + +def has_request_transforms() -> bool: + from litellm.proxy.hooks import PROXY_HOOKS + + builtins: Final = frozenset(PROXY_HOOKS.values()) + hooks: Final = ("async_pre_call_hook", "async_pre_request_hook", "async_pre_call_deployment_hook") + callbacks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomLogger) + return any( + type(callback) not in builtins + and any(getattr(type(callback), hook) is not getattr(CustomLogger, hook) for hook in hooks) + for callback in callbacks + ) + + +def _buckets(prefix_tokens: int, suffix_tokens: int, read_tokens: int, ttl_seconds: int) -> CacheTokenBuckets: + return CacheTokenBuckets( + uncached_input_tokens=suffix_tokens, + cache_read_input_tokens=read_tokens, + cache_creation_5m_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 300 else 0, + cache_creation_1h_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 3600 else 0, + ) + + +def _scenario(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> CacheCostScenario | None: + cost: Final = price_cache_tokens(model=model, deployment_id=deployment_id, tokens=tokens) + return CacheCostScenario(tokens=tokens, input_cost=cost) if cost is not None else None + + +def _capacity_counter( + limiter: _PROXY_MaxParallelRequestsHandler_v3, + caller: UserAPIKeyAuth, + model_name: str, + request_data: Mapping[str, object], +) -> TokenCounter: + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + async with limiter.request_capacity(caller, model_name, request_data=request_data): + return await count_prompt_tokens(model, api_key, body) + + return count + + +def _capacity_request_data( + http_request: Request, caller: UserAPIKeyAuth, request_data: Mapping[str, object] +) -> Mapping[str, object]: + # The parsed-body cache retains only original top-level keys. Replay the + # shared idempotent tag merges on limiter-only data when auth added metadata. + data: Final = dict(request_data) # mutable-ok: the existing tag merge owners accept a dictionary out-param + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(http_request, data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner takes the validated capacity dictionary + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner merges trusted key tags into capacity metadata + return MappingProxyType(data) + + +async def predict_arm( + deployment: Deployment, + body: Mapping[str, JsonValue], + prefix: PromptPrefix, + caller_key_hash: str, + cache: DualCache, + token_counter: TokenCounter, +) -> CachePredictionArm: + deployment_id: Final = deployment.model_info.id or "" + params: Final = deployment.litellm_params + unknown: Final = CachePredictionArm(deployment_id=deployment_id, model=params.model) + if deployment.model_info.blocked: + return unknown.model_copy(update=MappingProxyType({"reason": "unsupported_deployment_configuration"})) + target: Final = resolve_prediction_target(params) + if isinstance(target, UnsupportedPredictionTarget): + return unknown.model_copy(update=MappingProxyType({"reason": target.reason})) + model: Final = target.model + api_key: Final = target.api_key + total_count: Final = await token_counter(model, api_key, body) + prefix_count: Final = await token_counter(model, api_key, prefix.prefix_body) + if total_count is None or prefix_count is None or total_count < prefix_count: + return unknown.model_copy(update=MappingProxyType({"reason": "token_count_unavailable"})) + scope: Final = cache_scope(caller_key_hash, deployment_id, api_key, model) + observation: Final = await lookup(cache, scope, prefix) + exact: Final = observation is not None and observation.fingerprint == prefix.fingerprint + cacheable: Final = observation.cached_tokens if exact and observation is not None else prefix_count + if cacheable > total_count or (observation is not None and observation.cached_tokens > cacheable): + return unknown.model_copy(update=MappingProxyType({"reason": "inconsistent_prefix_token_count"})) + suffix: Final = total_count - cacheable + evidence: Final = ( + CacheEvidence(observed_at=observation.observed_at, expires_at=observation.expires_at) + if observation is not None + else None + ) + if cacheable < get_prompt_cache_min_tokens(params.model): + disabled: Final = _scenario(model, deployment_id, CacheTokenBuckets(uncached_input_tokens=total_count)) + if disabled is None: + return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"})) + return CachePredictionArm( + deployment_id=deployment_id, + model=model, + cache_state="disabled", + reason="below_cache_minimum", + estimate=disabled, + cold=disabled, + warm=disabled, + token_count_source="anthropic_count_tokens", + ) + fresh: Final = observation is not None and observation.expires_at > time.time() + read: Final = observation.cached_tokens if fresh and observation is not None else 0 + with pinned_billing_time(current_billing_time()): + cold: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, 0, prefix.ttl_seconds)) + warm: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, cacheable, prefix.ttl_seconds)) + estimate: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, read, prefix.ttl_seconds)) + if cold is None or warm is None or estimate is None: + return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"})) + return CachePredictionArm( + deployment_id=deployment_id, + model=model, + cache_state="warm" if fresh and exact else "partial" if fresh else "stale" if observation else "unknown", + reason=None if fresh else "observation_expired" if observation else "no_compatible_observation", + estimate=estimate, + cold=cold, + warm=warm, + evidence=evidence, + token_count_source="anthropic_count_tokens", + ) + + +@router.post( + "/cost/predict-cache", + tags=["Cost Tracking"], # mutable-ok: FastAPI requires a list for OpenAPI tags + response_model=CachePredictionResponse, +) +async def predict_cache_cost( + request: CachePredictionRequest, + http_request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CachePredictionResponse: + """Compare the next native Anthropic request on two configured deployment IDs. + + Estimates use provider token counting and recent successful cache telemetry for this key. + Unknown cache state uses the cold scenario when prices/counts are available. Cache observations + do not guarantee retention. v0 supports one message-content breakpoint, text and client tools; + system/tool-only breakpoints, thinking, images, nondefault Anthropic versions, beta headers and + request transforms are unknown. + Each provider count consumes one RPM unit and holds concurrency capacity; a comparison uses + up to four counts. The legacy rate limiter returns unknown without contacting the provider. + This endpoint does not generate tokens, prewarm caches, choose a model or alter routing. + """ + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + if llm_router is None: + raise HTTPException(status_code=503, detail="Model router is unavailable") + deployments: Final = get_cache_prediction_deployments( + current_deployment_id=request.current_deployment_id, + candidate_deployment_id=request.candidate_deployment_id, + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if deployments is None: + raise HTTPException(status_code=404, detail="Deployment not found") + current, candidate = deployments + for deployment in (current, candidate): + await can_key_call_resolved_model( + model=deployment.model_name, + llm_model_list=llm_router.get_model_list(), + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + prefix: Final = parse_prompt(request.request) + caller: Final = user_api_key_dict.api_key + caller_settings: Final = _CallerSettings.model_validate(user_api_key_dict, from_attributes=True) + unsupported_transform: Final = bool(caller_settings.config) or has_request_transforms() + unsupported_headers: Final = not supported_prediction_headers(http_request.headers) + limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + if ( + prefix is None + or not caller + or unsupported_transform + or unsupported_headers + or not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3) + ): + reason: Final = ( + "unsupported_provider_headers" + if unsupported_headers + else "unsupported_request_transform" + if unsupported_transform + else "unsupported_prompt_shape" + if prefix is None + else "caller_identity_unavailable" + if not caller + else "limiter_unavailable" + ) + return CachePredictionResponse( + stay=CachePredictionArm(deployment_id=request.current_deployment_id, reason=reason), + switch=CachePredictionArm(deployment_id=request.candidate_deployment_id, reason=reason), + switch_delta=None, + cache_rebuild_penalty=None, + ) + request_data: Final = _capacity_request_data( + http_request, user_api_key_dict, _REQUEST_DATA.validate_python(await _read_request_body(http_request)) + ) + stay: Final = await predict_arm( + current, + request.request, + prefix, + caller, + proxy_logging_obj.internal_usage_cache.dual_cache, + _capacity_counter(limiter, user_api_key_dict, current.model_name, request_data), + ) + switch: Final = ( + stay + if current.model_info.id == candidate.model_info.id + else await predict_arm( + candidate, + request.request, + prefix, + caller, + proxy_logging_obj.internal_usage_cache.dual_cache, + _capacity_counter(limiter, user_api_key_dict, candidate.model_name, request_data), + ) + ) + return CachePredictionResponse( + stay=stay, + switch=switch, + switch_delta=(switch.estimate.input_cost - stay.estimate.input_cost) + if switch.estimate is not None and stay.estimate is not None + else None, + cache_rebuild_penalty=(switch.estimate.input_cost - switch.warm.input_cost) + if switch.estimate is not None and switch.warm is not None + else None, + ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 4be0235adbb..b310fc661c4 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -270,6 +270,24 @@ class PassThroughStreamingHandler: - Vertex AI - OpenAI """ + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + _is_message_stop_chunk, # pyright: ignore[reportPrivateUsage] # both native stream paths share terminal-event detection + _is_provider_error_chunk, # pyright: ignore[reportPrivateUsage] # provider errors must not become cache evidence + ) + + # Transport reads can split event names and JSON payloads. Recognize terminal + # events only after the shared SSE framer has reassembled the collected bytes. + complete_frames, incomplete_tail = split_complete_sse_frames( + b"".join(raw_bytes) if endpoint_type == EndpointType.ANTHROPIC else b"" + ) + litellm_logging_obj.model_call_details[ # rebind-ok: stamp evidence on the per-request state read by callbacks + "prompt_cache_response_complete" + ] = ( + endpoint_type == EndpointType.ANTHROPIC + and not incomplete_tail.strip() + and _is_message_stop_chunk(complete_frames) + and not _is_provider_error_chunk(complete_frames) + ) try: ( standard_logging_response_object, diff --git a/litellm/types/management_endpoints/prompt_cache_prediction.py b/litellm/types/management_endpoints/prompt_cache_prediction.py new file mode 100644 index 00000000000..3789607b021 --- /dev/null +++ b/litellm/types/management_endpoints/prompt_cache_prediction.py @@ -0,0 +1,67 @@ +from collections.abc import Mapping +from typing import Annotated, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt + +TokenCount: TypeAlias = Annotated[StrictInt, Field(ge=0)] + + +class CacheTokenBuckets(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + uncached_input_tokens: TokenCount = 0 + cache_read_input_tokens: TokenCount = 0 + cache_creation_5m_input_tokens: TokenCount = 0 + cache_creation_1h_input_tokens: TokenCount = 0 + + @property + def total_tokens(self) -> int: + return ( + self.uncached_input_tokens + + self.cache_read_input_tokens + + self.cache_creation_5m_input_tokens + + self.cache_creation_1h_input_tokens + ) + + +class CacheEvidence(BaseModel): + model_config = ConfigDict(frozen=True) + + observed_at: float + expires_at: float + source: Literal["provider_usage"] = "provider_usage" + confidence: Literal["observed"] = "observed" + + +class CacheCostScenario(BaseModel): + tokens: CacheTokenBuckets + input_cost: float + + +class CachePredictionArm(BaseModel): + deployment_id: str + model: str | None = None + cache_state: Literal["warm", "partial", "stale", "unknown", "disabled"] = "unknown" + reason: str | None = None + estimate: CacheCostScenario | None = None + cold: CacheCostScenario | None = None + warm: CacheCostScenario | None = None + evidence: CacheEvidence | None = None + token_count_source: Literal["anthropic_count_tokens"] | None = None + + +class CachePredictionRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + + current_deployment_id: str = Field(min_length=1, max_length=256) + candidate_deployment_id: str = Field(min_length=1, max_length=256) + request: Mapping[str, JsonValue] + + +class CachePredictionResponse(BaseModel): + stay: CachePredictionArm + switch: CachePredictionArm + switch_delta: float | None + cache_rebuild_penalty: float | None + pricing_basis: Literal["input_before_discounts_and_margins"] = "input_before_discounts_and_margins" + cache_guarantee: Literal[False] = False diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py new file mode 100644 index 00000000000..2b36866a1a0 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -0,0 +1,209 @@ +import json +from collections.abc import Mapping +from datetime import datetime +from types import SimpleNamespace +from typing import Final + +import httpx +import pytest +import respx +from pydantic import JsonValue + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.anthropic.count_tokens import handler as count_handler +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION +from litellm.llms.anthropic.prompt_cache_prediction import ( + NativePredictionTarget, + cache_scope, + count_prompt_tokens, + parse_observed_cache, + parse_prompt, + resolve_prediction_target, + supported_prediction_headers, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.models.credentials import CredentialItem +from litellm.proxy import proxy_server +from litellm.proxy.hooks.prompt_cache_prediction import PromptCacheObserver, lookup +from litellm.proxy.management_endpoints.prompt_cache_prediction import predict_arm +from litellm.proxy.utils import InternalUsageCache +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage + +_MODEL: Final = "claude-sonnet-5" +_KEY: Final = "test-provider-key" +_CALLER: Final = "test-caller-hash" +_DEPLOYMENT: Final = "test-native-deployment" + + +def _body() -> dict[str, JsonValue]: + return { + "model": _MODEL, + "system": "Keep this context", + "tools": [{"name": "lookup", "input_schema": {"type": "object"}}], + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "A cacheable prefix", "cache_control": {"type": "ephemeral"}} + ]}], + } + + +@pytest.mark.parametrize("version", [None, "2099-01-01", DEFAULT_ANTHROPIC_API_VERSION]) +@pytest.mark.asyncio +async def test_observer_records_only_version_supported_by_token_counter(version: str | None) -> None: + cache: Final = DualCache() + observer: Final = PromptCacheObserver(InternalUsageCache(dual_cache=cache), clock=lambda: 1010.0) + body: Final = _body() + prefix: Final = parse_prompt(body) + assert prefix is not None + headers: Final = {"x-api-key": _KEY, **({"anthropic-version": version} if version is not None else {})} + wire: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages", headers=headers, json=body) + response: Final = ModelResponse( + model=_MODEL, + usage=Usage( + prompt_tokens=311, + completion_tokens=2, + total_tokens=313, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=100, + cache_creation_tokens=200, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_5m_input_tokens=200, ephemeral_1h_input_tokens=0 + ), + ), + ), + ) + await observer.async_log_success_event( + { + "call_type": "anthropic_messages", + "custom_llm_provider": "anthropic", + "httpx_response": httpx.Response(200, request=wire), + "first_api_call_start_time": datetime.fromtimestamp(1000.0), + "standard_logging_object": { + "status": "success", "model_id": _DEPLOYMENT, + "metadata": {"user_api_key_hash": _CALLER}, + }, + }, + response, + datetime.fromtimestamp(1010.0), + datetime.fromtimestamp(1010.0), + ) + default_scope: Final = cache_scope(_CALLER, _DEPLOYMENT, _KEY, _MODEL) + found: Final = await lookup(cache, default_scope, prefix, now=1010.0) + assert (found is not None) == (version == DEFAULT_ANTHROPIC_API_VERSION) + if version != DEFAULT_ANTHROPIC_API_VERSION: + other_scope: Final = cache_scope(_CALLER, _DEPLOYMENT, _KEY, _MODEL, version or "") + assert await lookup(cache, other_scope, prefix, now=1010.0) is None + + +@pytest.mark.parametrize("headers, supported", [ + ({}, True), + ({"Anthropic-Version": DEFAULT_ANTHROPIC_API_VERSION}, True), + ({"anthropic-version": "2099-01-01"}, False), + ({"Anthropic-Beta": ""}, False), + ({"anthropic-beta": "future-feature"}, False), +]) +def test_prediction_header_eligibility(headers: Mapping[str, str], supported: bool) -> None: + assert supported_prediction_headers(headers) is supported + + +@pytest.mark.asyncio +async def test_provider_count_uses_same_version_and_preserves_native_input(monkeypatch: pytest.MonkeyPatch) -> None: + body: Final = _body() + requests: Final[list[httpx.Request]] = [] + + def provider(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"input_tokens": 311}) + + client: Final = AsyncHTTPHandler() + await client.client.aclose() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + monkeypatch.setattr(count_handler, "get_async_httpx_client", lambda **kwargs: client) + try: + assert await count_prompt_tokens(_MODEL, _KEY, body) == 311 + finally: + await client.client.aclose() + assert len(requests) == 1 + assert requests[0].headers["anthropic-version"] == DEFAULT_ANTHROPIC_API_VERSION + assert requests[0].url == "https://api.anthropic.com/v1/messages/count_tokens" + assert json.loads(requests[0].content) == body + + +@pytest.mark.parametrize("source", ["static", "database"]) +@pytest.mark.asyncio +async def test_environment_credential_matches_native_count_and_observed_scope( + source: str, monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LIT7658_PROVIDER_KEY", _KEY) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + params: Final = { + "model": f"anthropic/{_MODEL}", "api_key": "os.environ/LIT7658_PROVIDER_KEY", + "api_base": "https://api.anthropic.com", + } + router: Final = litellm.Router(model_list=[{ + "model_name": "test-native", "litellm_params": dict(params), "model_info": {"id": _DEPLOYMENT}, + }] if source == "static" else [], num_retries=0) + if source == "database": + monkeypatch.setattr(proxy_server, "llm_router", router) + assert proxy_server.ProxyConfig()._add_deployment([SimpleNamespace( + model_id=_DEPLOYMENT, model_name="test-native", model_info={}, litellm_params=dict(params), + )]) == 1 + deployment: Final = router.get_deployment(_DEPLOYMENT) + assert deployment is not None + target: Final = resolve_prediction_target(deployment.litellm_params) + assert isinstance(target, NativePredictionTarget) + body: Final = _body() + with respx.mock() as upstream: + native: Final = upstream.post("https://api.anthropic.com/v1/messages").respond(200, json={ + "id": "msg_test", "type": "message", "role": "assistant", "model": _MODEL, + "content": [{"type": "text", "text": "Hello"}], "stop_reason": "end_turn", "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 1, "cache_read_input_tokens": 300}, + }) + counter: Final = upstream.post("https://api.anthropic.com/v1/messages/count_tokens").respond( + 200, json={"input_tokens": 311}, + ) + await router.aanthropic_messages( + model="test-native", max_tokens=1, **{key: value for key, value in body.items() if key != "model"}, + ) + assert await count_prompt_tokens(target.model, target.api_key, body) == 311 + assert native.call_count == counter.call_count == 1 + assert native.calls.last.request.headers["x-api-key"] == counter.calls.last.request.headers["x-api-key"] == _KEY + observed: Final = parse_observed_cache(native.calls.last.request, ModelResponse( + model=_MODEL, usage=Usage( + prompt_tokens=311, completion_tokens=1, total_tokens=312, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300), + ), + ), _CALLER, _DEPLOYMENT) + assert observed is not None + assert observed.scope == cache_scope(_CALLER, _DEPLOYMENT, target.api_key, target.model) + + +@pytest.mark.parametrize("inline_key", [None, _KEY]) +@pytest.mark.asyncio +async def test_named_credential_is_explicitly_unsupported_before_count( + inline_key: str | None, monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "credential_list", [CredentialItem( + credential_name="test-named", credential_info={}, credential_values={"api_key": "test-named-provider-key"}, + )]) + deployment: Final = Deployment( + model_name="test-native", + litellm_params=LiteLLM_Params( + model=f"anthropic/{_MODEL}", api_key=inline_key, litellm_credential_name="test-named", + ), + model_info=ModelInfo(id=_DEPLOYMENT), + ) + body: Final = _body() + prefix: Final = parse_prompt(body) + assert prefix is not None + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + pytest.fail("Unsupported named credentials must not reach provider counting") + + arm: Final = await predict_arm(deployment, body, prefix, _CALLER, DualCache(), count) + assert arm.cache_state == "unknown" + assert arm.reason == "unsupported_deployment_configuration" + assert arm.estimate is None and arm.cold is None and arm.warm is None diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index cdf1f897707..bd6a14cad21 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -433,6 +433,232 @@ def test_get_model_from_request_no_request_extracts_model(): ) +def _cache_prediction_router(): + from litellm.router import Router + + return Router(model_list=[ + { + "model_name": group, + "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"}, + "model_info": {"id": deployment_id, "team_id": team_id}, + } + for group, deployment_id, team_id in ( + ("current-group", "current-id", None), ("candidate-group", "candidate-id", None), + ("own-group", "own-id", "prediction-team"), ("foreign-group", "foreign-id", "foreign-team"), + ) + ]) + + +@pytest.mark.parametrize("candidate,team_id,expected", [ + ("candidate-id", None, ["current-group", "candidate-group"]), + ("current-id", None, "current-group"), + ("missing-id", None, None), + ("candidate-group", None, None), + ("own-id", None, None), + ("own-id", "prediction-team", ["current-group", "own-group"]), + ("foreign-id", "prediction-team", None), +]) +def test_cache_prediction_auth_resolves_only_exact_deployment_ids(candidate, team_id, expected): + assert get_model_from_request( + request_data={ + "current_deployment_id": "current-id", "candidate_deployment_id": candidate, + "request": {"model": "caller-controlled-provider-model"}, + }, + route="/cost/predict-cache", + llm_router=_cache_prediction_router(), + team_id=team_id, + ) == expected + + +def _cache_prediction_auth_app( + monkeypatch, allowed_routes, user_models, metadata=None, *, team_id=None, key_models=None, team_models=None +): + import importlib + from unittest.mock import AsyncMock + + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, ProxyException + from litellm.proxy.auth import auth_checks + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.management_endpoints import prompt_cache_prediction as endpoint + from litellm.proxy.utils import InternalUsageCache, ProxyLogging + + auth = importlib.import_module("litellm.proxy.auth.user_api_key_auth") + router = _cache_prediction_router() + allowed_models = ["current-group", "candidate-group", "own-group"] + token = UserAPIKeyAuth( + api_key="test-proxy-key-hash", user_id="prediction-user", user_role=LitellmUserRoles.INTERNAL_USER, + models=allowed_models if key_models is None else key_models, team_id=team_id, + team_models=allowed_models if team_models is None else team_models, + allowed_routes=allowed_routes, metadata=metadata or {}, + ) + user = LiteLLM_UserTable( + user_id=token.user_id, user_role=LitellmUserRoles.INTERNAL_USER.value, models=user_models, + ) + async def authenticate(request, request_data, **_headers): + await auth._enforce_key_and_fallback_model_access( + valid_token=token, request_data=request_data, route=request.url.path, request=request, + llm_model_list=router.get_model_list(), llm_router=router, + ) + return token + + monkeypatch.setattr(auth, "_user_api_key_auth_builder", authenticate) + monkeypatch.setattr(auth, "get_user_object", AsyncMock(return_value=user)) + team = LiteLLM_TeamTableCachedObj(team_id=team_id, models=token.team_models) if team_id else None + monkeypatch.setattr(auth, "get_team_object", AsyncMock(return_value=team)) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth, "get_global_proxy_spend", AsyncMock(return_value=0)) + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + monkeypatch.setattr(proxy_server, "user_custom_auth", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + logging = ProxyLogging(user_api_key_cache=DualCache()) + logging.proxy_hook_mapping["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler_v3( + InternalUsageCache(dual_cache=DualCache()) + ) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging) + counts = AsyncMock(return_value=6_000) + monkeypatch.setattr(endpoint, "count_prompt_tokens", counts) + app = FastAPI() + app.include_router(endpoint.router) + app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler) + return app, counts + + +def _cache_prediction_payload(candidate="candidate-id", current="current-id"): + return { + "current_deployment_id": current, "candidate_deployment_id": candidate, + "request": {"messages": [{"role": "user", "content": [{ + "type": "text", "text": "Stable cached context", + "cache_control": {"type": "ephemeral"}, + }]}]}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed_routes,user_models,candidate,status_code", [ + (["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403), + (["/cost/predict-cache"], ["current-group"], "candidate-id", 403), + (["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200), + (["/cost/predict-cache"], ["current-group"], "current-id", 200), + (["/cost/predict-cache"], ["current-group"], "missing-id", 404), +]) +async def test_cache_prediction_authorizes_route_and_personal_models_before_provider_counts( + monkeypatch, allowed_routes, user_models, candidate, status_code +): + import httpx + + app, counts = _cache_prediction_auth_app(monkeypatch, allowed_routes, user_models) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response = await client.post("/cost/predict-cache", json=_cache_prediction_payload(candidate)) + + assert response.status_code == status_code, response.text + if status_code == 200: + assert counts.await_count == (2 if candidate == "current-id" else 4) + else: + assert counts.await_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) +@pytest.mark.parametrize("team_id,key_models,user_models,team_models", [ + (None, ["*"], ["*"], None), + (None, ["current-group", "candidate-group"], ["*"], None), + (None, ["*"], ["current-group", "candidate-group"], None), + ("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]), +]) +async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_authorization( + monkeypatch, arm, team_id, key_models, user_models, team_models +): + import httpx + + app, counts = _cache_prediction_auth_app( + monkeypatch, ["/cost/predict-cache"], user_models, + team_id=team_id, key_models=key_models, team_models=team_models, + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + missing = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "missing-id"}) + foreign = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "foreign-id"}) + + assert missing.status_code == foreign.status_code == 404, foreign.text + assert missing.json() == foreign.json() == {"detail": "Deployment not found"} + assert counts.await_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) +@pytest.mark.parametrize("key_models,team_models,status_code", [ + (["*"], ["*"], 200), + (["current-group", "candidate-group"], ["*"], 403), + (["*"], ["current-group", "candidate-group"], 403), +]) +async def test_cache_prediction_checks_each_visible_team_deployment_model( + monkeypatch, arm, key_models, team_models, status_code +): + import httpx + + app, counts = _cache_prediction_auth_app( + monkeypatch, ["/cost/predict-cache"], ["*"], + team_id="prediction-team", key_models=key_models, team_models=team_models, + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "own-id"}) + + assert response.status_code == status_code, response.text + assert counts.await_count == (4 if status_code == 200 else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) +async def test_cache_prediction_checks_each_visible_personal_deployment_model(monkeypatch, arm): + import httpx + + app, counts = _cache_prediction_auth_app(monkeypatch, ["/cost/predict-cache"], ["current-group"]) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response = await client.post( + "/cost/predict-cache", json={**_cache_prediction_payload(candidate="current-id"), arm: "candidate-id"} + ) + + assert response.status_code == 403, response.text + assert response.json()["error"]["type"] == "user_model_access_denied" + assert counts.await_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("header_tag,key_tags,limit,status_code,provider_calls", [ + ("limited", [], 1, 429, 1), + (None, ["limited"], 1, 429, 1), + ("limited", ["limited"], 4, 200, 4), + ("unlimited", [], 1, 200, 4), +]) +async def test_cache_prediction_preserves_authenticated_header_and_key_tag_rpm( + monkeypatch, header_tag, key_tags, limit, status_code, provider_calls +): + import httpx + + app, counts = _cache_prediction_auth_app( + monkeypatch, ["/cost/predict-cache"], ["current-group", "candidate-group"], + metadata={"tag_rpm_limit": {"limited": limit}, "tags": key_tags}, + ) + headers = {"x-litellm-tags": header_tag} if header_tag else {} + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response = await client.post("/cost/predict-cache", json=_cache_prediction_payload(), headers=headers) + assert response.status_code == status_code, response.text + assert counts.await_count == provider_calls + if limit == 4: + exhausted = await client.post("/cost/predict-cache", json=_cache_prediction_payload(), headers=headers) + assert exhausted.status_code == 429, exhausted.text + assert counts.await_count == 4 + assert all("metadata" not in call.args[2] for call in counts.await_args_list) + + def test_get_model_from_request_supports_google_model_names_with_slashes(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py b/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py new file mode 100644 index 00000000000..994684a6005 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py @@ -0,0 +1,105 @@ +from typing import Final + +import pytest + +import litellm +from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens +from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets + + +@pytest.mark.parametrize( + ("model", "expected"), + [("anthropic/claude-sonnet-4-5", 1.26), ("anthropic/claude-sonnet-4-6", 0.63)], +) +def test_prices_all_cache_buckets_at_total_context_tier(model: str, expected: float) -> None: + tokens: Final = CacheTokenBuckets( + uncached_input_tokens=100_000, + cache_read_input_tokens=50_000, + cache_creation_5m_input_tokens=20_000, + cache_creation_1h_input_tokens=40_000, + ) + assert price_cache_tokens(model, "unconfigured-deployment", tokens) == pytest.approx(expected) + + +@pytest.mark.parametrize(("total", "expected"), [(200_000, 0.387), (200_001, 0.774006)]) +def test_long_context_tier_starts_above_threshold(total: int, expected: float) -> None: + tokens: Final = CacheTokenBuckets( + uncached_input_tokens=total - 100_000, + cache_creation_1h_input_tokens=10_000, + cache_read_input_tokens=90_000, + ) + actual: Final = price_cache_tokens("anthropic/claude-sonnet-4-5", "unconfigured-deployment", tokens) + assert actual == pytest.approx(expected) + + +def test_deployment_tariff_wins_without_proxy_discounts_or_margins(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost.copy()) + litellm.Router( + model_list=[ + { + "model_name": "cache-pricing-test", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-6", + "api_key": "test-only", + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "cache_read_input_token_cost": 0.000001, + "cache_creation_input_token_cost": 0.0000125, + "cache_creation_input_token_cost_above_1hr": 0.00002, + }, + "model_info": {"id": "cache-pricing-test-a"}, + } + ] + ) + monkeypatch.setattr(litellm, "cost_discount_config", {"anthropic": 0.5}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"percentage": 0.3, "fixed_amount": 1.0}}) + tokens: Final = CacheTokenBuckets( + uncached_input_tokens=3_000, + cache_read_input_tokens=4_000, + cache_creation_5m_input_tokens=1_000, + cache_creation_1h_input_tokens=2_000, + ) + assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-test-a", tokens) == pytest.approx(0.0865) + + +@pytest.mark.parametrize("rate", [None, -1.0, float("nan"), float("inf"), "0.00001", True]) +def test_unknown_for_absent_or_invalid_active_cache_rate(monkeypatch: pytest.MonkeyPatch, rate: object) -> None: + monkeypatch.setitem( + litellm.model_cost, + "cache-pricing-invalid", + { + "litellm_provider": "anthropic", + "mode": "chat", + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "cache_creation_input_token_cost_above_1hr": rate, + }, + ) + tokens: Final = CacheTokenBuckets(cache_creation_1h_input_tokens=4_000) + assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-invalid", tokens) is None + + +def test_missing_input_price_is_unknown_even_when_get_model_info_defaults_to_zero( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(litellm.model_cost, "cache-pricing-missing", {"litellm_provider": "anthropic", "mode": "chat"}) + tokens: Final = CacheTokenBuckets(uncached_input_tokens=4_000) + assert price_cache_tokens("cache-pricing-missing", "unconfigured-deployment", tokens) is None + + +def test_explicit_free_pricing_is_not_unknown(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem( + litellm.model_cost, + "cache-pricing-free", + { + "litellm_provider": "anthropic", + "mode": "chat", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "cache_read_input_token_cost": 0.0, + "cache_creation_input_token_cost": 0.0, + "cache_creation_input_token_cost_above_1hr": 0.0, + }, + ) + tokens: Final = CacheTokenBuckets(uncached_input_tokens=100, cache_read_input_tokens=5_000) + assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-free", tokens) == 0.0 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 10c0bb88a82..48f980086fd 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6284,3 +6284,248 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_ assert isinstance(values, list) assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) + + +@pytest.mark.parametrize( + "limits, request_data, counter_scope", + [ + ({"rpm_limit": 1}, {}, "api_key"), + ({"user_id": "u", "user_rpm_limit": 1}, {}, "user"), + ({"team_id": "t", "team_rpm_limit": 1}, {}, "team"), + ( + {"team_id": "t", "user_id": "u", "team_member_rpm_limit": 1}, + {}, + "team_member", + ), + ({"end_user_id": "e", "end_user_rpm_limit": 1}, {}, "end_user"), + ( + {"metadata": {"model_rpm_limit": {"test-model": 1}}}, + {}, + "model_per_key", + ), + ( + {"metadata": {"tag_rpm_limit": {"test-tag": 1}}}, + {"metadata": {"tags": ["test-tag"]}}, + "tag_per_key", + ), + ( + { + "team_id": "t", + "metadata": {"model_rpm_limit": {"test-model": 100}}, + "team_metadata": {"model_rpm_limit": {"test-model": 1}}, + }, + {}, + "model_per_team", + ), + ( + {"project_id": "p", "project_metadata": {"model_rpm_limit": {"test-model": 1}}}, + {}, + "model_per_project", + ), + ({"org_id": "o", "organization_rpm_limit": 1}, {}, "organization"), + ( + {"org_id": "o", "organization_metadata": {"model_rpm_limit": {"test-model": 1}}}, + {}, + "model_per_organization", + ), + ], +) +@pytest.mark.parametrize("request_kind", ["count", "generation"]) +@pytest.mark.asyncio +async def test_request_capacity_enforces_shared_rpm_scopes( + limits, request_data, counter_scope, request_kind +): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth = UserAPIKeyAuth(api_key=hash_token("sk-count-rpm"), **limits) + async def request(): + if request_kind == "generation": + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={**request_data, "model": "test-model"}, + call_type="acompletion", + ) + return + async with handler.request_capacity(auth, "test-model", request_data=request_data): + pass + + await request() + with pytest.raises(HTTPException) as exc: + await request() + assert exc.value.status_code == 429 + assert counter_scope in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + router = Router(model_list=[{ + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, + }]) + monkeypatch.setattr(proxy_server, "llm_router", router) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + auth = UserAPIKeyAuth( + api_key=hash_token("sk-count-dynamic"), + rpm_limit=1, + metadata={"rpm_limit_type": "dynamic"}, + ) + for _ in range(2): + async with handler.request_capacity(auth, "test-model"): + pass + router.cache.set_cache("test-deployment:fails", 100, ttl=60, local_only=True) + async with handler.request_capacity(auth, "test-model"): + pass + with pytest.raises(HTTPException) as exc: + async with handler.request_capacity(auth, "test-model"): + pytest.fail("dynamic RPM must enforce after deployment failures") + assert exc.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_request_capacity_skips_tokens_and_preserves_parent_stash(): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth = UserAPIKeyAuth( + api_key=hash_token("sk-count-tpm"), + rpm_limit=5, + tpm_limit=1, + max_parallel_requests=1, + project_id="p", + project_metadata={ + "model_tpm_limit": {"test-model": 1}, + "model_itpm_limit": {"test-model": 1}, + "model_otpm_limit": {"test-model": 1}, + }, + ) + token_scopes = ( + ("api_key", auth.api_key), + ("model_per_project", "p:test-model"), + ("model_per_project_itpm", "p:test-model"), + ("model_per_project_otpm", "p:test-model"), + ) + for scope, value in token_scopes: + token_key = handler.create_rate_limit_keys(scope, value, "tokens") + await cache.async_set_cache(token_key, 100, ttl=60) + await cache.async_set_cache(f"{{{scope}:{value}}}:window", int(time.time()), ttl=60) + parent = get_or_create_request_stash() + parent.reserved_tokens = 123 + parent.parallel_slot = ParallelSlotAcquisition(slot_id="parent", counter_keys=["parent-gauge"]) + for _ in range(2): + async with handler.request_capacity(auth, "test-model"): + assert get_request_stash() is parent + assert parent.parallel_slot["slot_id"] == "parent" + assert parent.reserved_tokens == 123 + for scope, value in token_scopes: + assert await cache.async_get_cache(handler.create_rate_limit_keys(scope, value, "tokens")) == 100 + + +@pytest.mark.parametrize("exit_mode", ["success", "failure", "cancel"]) +@pytest.mark.asyncio +async def test_request_capacity_releases_exact_parallel_slot(exit_mode): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth = UserAPIKeyAuth(api_key=hash_token("sk-count-parallel"), max_parallel_requests=1) + entered = asyncio.Event() + finish = asyncio.Event() + + async def provider(): + async with handler.request_capacity(auth, "test-model"): + entered.set() + await finish.wait() + if exit_mode == "failure": + raise RuntimeError("provider failed") + + task = asyncio.create_task(provider()) + await asyncio.wait_for(entered.wait(), timeout=2) + try: + for _ in range(2): + with pytest.raises(HTTPException) as exc: + async with handler.request_capacity(auth, "test-model"): + pytest.fail("rejected request freed the occupied slot") + assert exc.value.status_code == 429 + finally: + if exit_mode == "cancel": + task.cancel() + else: + finish.set() + if exit_mode == "success": + await task + else: + with pytest.raises(asyncio.CancelledError if exit_mode == "cancel" else RuntimeError): + await task + async with handler.request_capacity(auth, "test-model"): + pass + + +class _DelayedCapacityUsageCache: + def __init__(self): + self.delegate = InternalUsageCache(DualCache()) + self.dual_cache = self.delegate.dual_cache + self.acquired = asyncio.Event() + self.finish_admission = asyncio.Event() + self.releasing = asyncio.Event() + self.finish_release = asyncio.Event() + + async def async_get_cache(self, *args, **kwargs): + return await self.delegate.async_get_cache(*args, **kwargs) + + async def async_batch_get_cache(self, *args, **kwargs): + return await self.delegate.async_batch_get_cache(*args, **kwargs) + + async def async_set_cache(self, key, value, **kwargs): + await self.delegate.async_set_cache(key=key, value=value, **kwargs) + if not key.endswith(":max_parallel_requests"): + return + if value: + self.acquired.set() + await self.finish_admission.wait() + else: + self.releasing.set() + await self.finish_release.wait() + + +@pytest.mark.asyncio +async def test_request_capacity_finishes_admission_and_release_despite_repeated_cancel(): + cache = _DelayedCapacityUsageCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + auth = UserAPIKeyAuth(api_key=hash_token("sk-count-cancel-admission"), max_parallel_requests=1) + + async def provider(): + async with handler.request_capacity(auth, "test-model"): + pytest.fail("cancelled admission entered provider body") + + task = asyncio.create_task(provider()) + await asyncio.wait_for(cache.acquired.wait(), timeout=2) + task.cancel() + await asyncio.sleep(0) + cache.finish_admission.set() + await asyncio.wait_for(cache.releasing.wait(), timeout=2) + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + cache.finish_release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=2) + async with handler.request_capacity(auth, "test-model"): + pass + + +@pytest.mark.asyncio +async def test_request_capacity_rejection_keeps_existing_redis_mirror(): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth = UserAPIKeyAuth(api_key=hash_token("sk-count-mirror"), max_parallel_requests=1) + counter_key = handler.create_rate_limit_keys("api_key", auth.api_key, "max_parallel_requests") + await cache.async_set_cache(counter_key, 1, ttl=60, local_only=True) + for _ in range(2): + with pytest.raises(HTTPException) as exc: + async with handler.request_capacity(auth, "test-model"): + pytest.fail("rejection released another request's mirrored slot") + assert exc.value.status_code == 429 + assert await cache.async_get_cache(counter_key, local_only=True) == 1 diff --git a/tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py b/tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py new file mode 100644 index 00000000000..82af3e9a6ef --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py @@ -0,0 +1,300 @@ +import asyncio +import json +import time +from datetime import datetime + +import httpx +import pytest + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.anthropic.prompt_cache_prediction import cache_scope, parse_prompt +from litellm.proxy.hooks.prompt_cache_prediction import ( + PromptCacheObserver, + lookup, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.types.utils import ModelResponse + +MODEL = "claude-sonnet-5" +CALLER = "a" * 64 +DEPLOYMENT = "native-deployment" +KEY = "test-provider-key" + + +def body(ttl="5m", texts=("private cache prefix",)): + return { + "model": MODEL, + "max_tokens": 2, + "system": "private system instructions", + "tools": [{"name": "lookup", "input_schema": {"type": "object"}}], + "messages": [{"role": "user", "content": [ + {"type": "text", "text": text, **( + {"cache_control": {"type": "ephemeral", "ttl": ttl}} + if index == len(texts) - 1 else {} + )} + for index, text in enumerate(texts) + ]}], + } + + +def usage(ttl="5m", read=100, write=200): + return { + "input_tokens": 11, + "output_tokens": 2, + "cache_read_input_tokens": read, + "cache_creation_input_tokens": write, + "cache_creation": { + "ephemeral_5m_input_tokens": write if ttl == "5m" else 0, + "ephemeral_1h_input_tokens": write if ttl == "1h" else 0, + }, + } + + +def event(request_body, started=1000.0, headers=None, **overrides): + request = httpx.Request( + "POST", "https://api.anthropic.com/v1/messages", json=request_body, + headers={"x-api-key": KEY, "anthropic-version": "2023-06-01", **(headers or {})}, + ) + return { + "call_type": "anthropic_messages", + "custom_llm_provider": "anthropic", + "cache_hit": False, + "httpx_response": httpx.Response(200, request=request), + "first_api_call_start_time": datetime.fromtimestamp(started), + "standard_logging_object": { + "status": "success", "model_id": DEPLOYMENT, + "metadata": {"user_api_key_hash": CALLER}, + }, + **overrides, + } + + +async def observe(cache, request_body=None, native_usage=None, now=1010.0, **overrides): + observer = PromptCacheObserver(InternalUsageCache(dual_cache=cache), clock=lambda: now) + response = ModelResponse( + model=MODEL, + usage=AnthropicConfig().calculate_usage(native_usage or usage(), reasoning_content=None), + ) + await observer.async_log_success_event( + event(request_body or body(), **overrides), response, + datetime.fromtimestamp(now), datetime.fromtimestamp(now), + ) + + +def scope(**overrides): + return cache_scope(**{ + "caller_key_hash": CALLER, "deployment_id": DEPLOYMENT, + "provider_key": KEY, "model": MODEL, **overrides, + }) + + +@pytest.mark.parametrize("ttl,expires", [("5m", 1300), ("1h", 4600)]) +@pytest.mark.asyncio +async def test_observed_cache_count_and_request_start_expiry_survive_as_stale(ttl, expires): + cache = DualCache() + request_body = body(ttl=ttl) + await observe(cache, request_body, usage(ttl=ttl)) + prefix = parse_prompt(request_body) + observed = await lookup(cache, scope(), prefix, now=1200) + assert observed.cached_tokens == 300 + assert observed.observed_at == 1010 + assert observed.expires_at == expires + assert await lookup(cache, scope(), prefix, now=expires) == observed + saved = json.dumps(cache.in_memory_cache.cache_dict) + assert "private cache prefix" not in saved + assert "private system instructions" not in saved + assert KEY not in saved + assert CALLER not in saved + + +@pytest.mark.parametrize("changed", [ + {"caller_key_hash": "b" * 64}, {"deployment_id": "other"}, + {"provider_key": "rotated"}, {"model": "claude-opus-5"}, + {"anthropic_version": "different"}, +]) +@pytest.mark.asyncio +async def test_cache_evidence_is_isolated_by_every_scope_dimension(changed): + cache = DualCache() + await observe(cache) + assert await lookup(cache, scope(**changed), parse_prompt(body()), now=1010) is None + + +@pytest.mark.asyncio +async def test_append_only_prefix_finds_prior_evidence_but_edit_or_context_change_does_not(): + cache = DualCache() + await observe(cache) + extended = parse_prompt(body(texts=("private cache prefix", "new turn"))) + prior = await lookup(cache, scope(), extended, now=1010) + assert prior.cached_tokens == 300 + assert prior.fingerprint != extended.fingerprint + for changed in ( + body(texts=("edited prefix", "new turn")), + {**body(), "system": "different system"}, + {**body(), "tools": [{"name": "other", "input_schema": {"type": "object"}}]}, + body(ttl="1h"), + ): + assert await lookup(cache, scope(), parse_prompt(changed), now=1010) is None + outside_lookback = parse_prompt(body(texts=("private cache prefix", *[str(i) for i in range(20)]))) + assert await lookup(cache, scope(), outside_lookback, now=1010) is None + + +@pytest.mark.parametrize("change", [ + {"thinking": {"type": "enabled", "budget_tokens": 1024}}, + {"tool_choice": {"type": "auto"}}, + {"cache_control": {"type": "ephemeral"}}, + {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + {"system": [{"type": "text", "text": "system", "cache_control": {"type": "ephemeral"}}]}, + {"messages": [{"role": "user", "content": [{"type": "image", "source": {}}]}]}, + {"messages": [{"role": "user", "content": "no breakpoint"}]}, +]) +def test_unsupported_or_ambiguous_shapes_have_no_cache_identity(change): + assert parse_prompt({**body(), **change}) is None + duplicate = body() + duplicate["messages"][0]["content"].append(duplicate["messages"][0]["content"][0]) + assert parse_prompt(duplicate) is None + + +@pytest.mark.parametrize("overrides", [ + {"cache_hit": True}, {"call_type": "completion"}, + {"custom_llm_provider": "bedrock"}, {"stream": True}, + {"headers": {"anthropic-beta": "unverified-feature"}}, + {"headers": {"x-custom-header": "unverified"}}, + {"standard_logging_object": {"status": "success", "model_id": DEPLOYMENT, "metadata": {}}}, +]) +@pytest.mark.asyncio +async def test_unverified_source_never_creates_observations(overrides): + cache = DualCache() + await observe(cache, **overrides) + assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None + + +@pytest.mark.parametrize("native_usage", [ + usage(write=0), + {**usage(), "cache_creation": None}, + {**usage(), "cache_creation": {"ephemeral_5m_input_tokens": 199, "ephemeral_1h_input_tokens": 0}}, + usage(ttl="1h"), + {**usage(), "cache_creation_input_tokens": -200}, +]) +@pytest.mark.asyncio +async def test_missing_or_contradictory_telemetry_cannot_create_observations(native_usage): + cache = DualCache() + await observe(cache, native_usage=native_usage) + assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None + + +@pytest.mark.asyncio +async def test_pure_read_refresh_requires_prior_matching_evidence(): + cache = DualCache() + await observe(cache, native_usage=usage(read=300, write=0)) + assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None + await observe(cache) + await observe(cache, native_usage=usage(read=300, write=0), started=1100, now=1110) + assert (await lookup(cache, scope(), parse_prompt(body()), now=1110)).expires_at == 1400 + + +class RecordingObserver(PromptCacheObserver): + def __init__(self, cache): + super().__init__(InternalUsageCache(dual_cache=cache)) + self.finished = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await super().async_log_success_event(kwargs, response_obj, start_time, end_time) + self.finished.set() + + +def native_response(): + return { + "id": "msg_prediction", "type": "message", "role": "assistant", "model": MODEL, + "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", + "stop_sequence": None, "usage": usage(ttl="1h"), + } + + +def stream_response(completed, provider_error=False): + response = native_response() + events = [ + {"type": "message_start", "message": {**response, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}}, + ] + if completed: + events.append({"type": "message_stop"}) + if provider_error: + events.append({"type": "error", "error": {"type": "overloaded_error", "message": "temporary failure"}}) + return "".join(f"event: {item['type']}\ndata: {json.dumps(item)}\n\n" for item in events) + + +class TransportChunks(httpx.AsyncByteStream): + def __init__(self, payload, chunk_size, fragment_error_only=False): + self.payload = payload.encode() + self.chunk_size = chunk_size or len(self.payload) + self.prefix_length = self.payload.index(b"event: error") if fragment_error_only else 0 + + async def __aiter__(self): + if self.prefix_length: + yield self.payload[:self.prefix_length] + for offset in range(self.prefix_length, len(self.payload), self.chunk_size): + yield self.payload[offset:offset + self.chunk_size] + + +@pytest.mark.parametrize("stream,completed,provider_error,transport", [ + (False, True, False, "whole"), + (True, True, False, "whole"), + (True, False, False, "whole"), + (True, True, True, "whole"), + (True, True, False, "fragmented"), + (True, False, False, "fragmented"), + (True, True, True, "fragmented"), + (True, True, True, "fragmented_error"), + (True, True, False, "unterminated"), +]) +@pytest.mark.asyncio +async def test_native_production_callback_records_only_completed_wire_requests(stream, completed, provider_error, transport): + cache = DualCache() + observer = RecordingObserver(cache) + litellm.logging_callback_manager.add_litellm_callback(observer) + + def provider(request): + if stream: + payload = stream_response(completed, provider_error) + if transport == "unterminated": + payload = payload.removesuffix("\n\n") + return httpx.Response( + 200, request=request, headers={"content-type": "text/event-stream"}, + stream=TransportChunks( + payload, 1 if transport.startswith("fragmented") else None, + fragment_error_only=transport == "fragmented_error", + ), + ) + return httpx.Response(200, request=request, json=native_response()) + + client = AsyncHTTPHandler() + await client.client.aclose() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + try: + request_body = body(ttl="1h") + before = time.time() + result = await litellm.anthropic_messages( + **{**request_body, "model": f"anthropic/{MODEL}"}, + api_key=KEY, client=client, stream=stream, model_info={"id": DEPLOYMENT}, + litellm_metadata={"user_api_key_hash": CALLER, "model_info": {"id": DEPLOYMENT}}, + ) + if stream: + async for _ in result: + pass + await asyncio.wait_for(observer.finished.wait(), timeout=5) + found = await lookup(cache, scope(), parse_prompt(request_body)) + if completed and not provider_error and transport != "unterminated": + assert found is not None + assert found.cached_tokens == 300 + assert before + 3600 <= found.expires_at <= time.time() + 3600 + else: + assert found is None + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(observer) + await client.client.aclose() diff --git a/tests/test_litellm/proxy/management_endpoints/test_prompt_cache_prediction.py b/tests/test_litellm/proxy/management_endpoints/test_prompt_cache_prediction.py new file mode 100644 index 00000000000..0ec277be884 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_prompt_cache_prediction.py @@ -0,0 +1,698 @@ +import asyncio +import time +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final, Literal + +import httpx +import pytest +from fastapi import FastAPI, Request +from pydantic import JsonValue + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.http_parsing_utils import _read_request_body, _safe_set_request_parsed_body +from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 +from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, cache_scope, parse_prompt +from litellm.proxy.hooks.prompt_cache_prediction import ( + CacheObservation, + _cache_key, +) +from litellm.proxy.management_endpoints import prompt_cache_prediction as endpoint +from litellm.proxy.utils import InternalUsageCache +from litellm.types.management_endpoints.prompt_cache_prediction import CachePredictionResponse +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + +_PROVIDER_KEY: Final = "cache-prediction-test-provider-key" +_CALLER: Final = "cache-prediction-test-caller-hash" + + +@pytest.fixture(autouse=True) +def anthropic_endpoint_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + + +def _body(ttl: str = "5m", *, extended: bool = False) -> dict[str, JsonValue]: + blocks: Final[list[JsonValue]] = [ + {"type": "text", "text": "Stable context"}, + *([{"type": "text", "text": "Appended context"}] if extended else []), + ] + return { + "max_tokens": 10, + "system": "Follow the project conventions", + "messages": [ + { + "role": "user", + "content": [ + *blocks[:-1], + {**blocks[-1], "cache_control": {"type": "ephemeral", "ttl": ttl}}, + {"type": "text", "text": "Follow-up question"}, + ], + } + ], + } + + +def _prefix(body: Mapping[str, JsonValue]) -> PromptPrefix: + prefix: Final = parse_prompt(body) + assert prefix is not None + return prefix + + +def _deployment( + deployment_id: str = "sonnet", + model: str = "claude-sonnet-5", + *, + team_id: str | None = None, + api_base: str | None = None, +) -> Deployment: + return Deployment( + model_name=deployment_id, + litellm_params=LiteLLM_Params(model=f"anthropic/{model}", api_key=_PROVIDER_KEY, api_base=api_base), + model_info=ModelInfo(id=deployment_id, team_id=team_id), + ) + + +@dataclass(frozen=True) +class Counts: + total: int | None = 6_000 + prefix: int | None = 5_000 + + async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + assert api_key == _PROVIDER_KEY + assert model.startswith("claude-") + return self.total if "max_tokens" in body else self.prefix + + +async def _observe( + cache: DualCache, + body: Mapping[str, JsonValue], + *, + deployment_id: str = "sonnet", + model: str = "claude-sonnet-5", + cached_tokens: int = 5_000, + expired: bool = False, + caller: str = _CALLER, +) -> None: + prefix: Final = _prefix(body) + now: Final = time.time() + observation: Final = CacheObservation( + fingerprint=prefix.fingerprint, + cached_tokens=cached_tokens, + observed_at=now - 400 if expired else now - 10, + expires_at=now - 100 if expired else now + 290, + ) + scope: Final = cache_scope(caller, deployment_id, _PROVIDER_KEY, model) + await cache.async_set_cache(_cache_key(scope, prefix.fingerprint), observation.model_dump_json(), ttl=3_600) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("ttl", "cold_cost"), [("5m", 0.0145), ("1h", 0.022)]) +async def test_unobserved_cache_prices_cold_and_warm_bounds(ttl: str, cold_cost: float) -> None: + body: Final = _body(ttl) + arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts()) + + assert arm.cache_state == "unknown" + assert arm.reason == "no_compatible_observation" + assert arm.evidence is None + assert arm.estimate is not None and arm.cold is not None and arm.warm is not None + assert arm.estimate.input_cost == pytest.approx(cold_cost) + assert arm.cold.input_cost == pytest.approx(cold_cost) + assert arm.warm.input_cost == pytest.approx(0.003) + assert arm.cold.tokens.uncached_input_tokens == 1_000 + assert arm.cold.tokens.cache_read_input_tokens == 0 + assert arm.cold.tokens.cache_creation_5m_input_tokens == (5_000 if ttl == "5m" else 0) + assert arm.cold.tokens.cache_creation_1h_input_tokens == (5_000 if ttl == "1h" else 0) + assert arm.warm.tokens.cache_read_input_tokens == 5_000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("cached_tokens", "warm_cost", "cold_cost"), [(5_400, 0.00228, 0.0147), (4_600, 0.00372, 0.0143)] +) +@pytest.mark.parametrize("expired", [False, True]) +async def test_exact_prefix_conserves_total_with_observed_count_in_all_scenarios( + cached_tokens: int, warm_cost: float, cold_cost: float, expired: bool +) -> None: + cache: Final = DualCache() + body: Final = _body() + await _observe(cache, body, cached_tokens=cached_tokens, expired=expired) + arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts()) + + assert arm.cache_state == ("stale" if expired else "warm") + assert arm.evidence is not None + assert arm.estimate is not None and arm.warm is not None and arm.cold is not None + assert arm.warm.tokens.cache_read_input_tokens == cached_tokens + assert arm.warm.tokens.cache_creation_5m_input_tokens == 0 + assert arm.cold.tokens.cache_creation_5m_input_tokens == cached_tokens + assert arm.cold.tokens.cache_read_input_tokens == 0 + for scenario in (arm.estimate, arm.cold, arm.warm): + assert scenario.tokens.total_tokens == 6_000 + assert scenario.tokens.uncached_input_tokens == 6_000 - cached_tokens + assert arm.warm.input_cost == pytest.approx(warm_cost) + assert arm.cold.input_cost == pytest.approx(cold_cost) + assert arm.estimate.input_cost == pytest.approx(cold_cost if expired else warm_cost) + + +@pytest.mark.asyncio +async def test_observed_prefix_larger_than_full_request_returns_unknown() -> None: + cache: Final = DualCache() + body: Final = _body() + await _observe(cache, body, cached_tokens=6_001) + arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts()) + + assert arm.cache_state == "unknown" + assert arm.reason == "inconsistent_prefix_token_count" + assert arm.estimate is None and arm.cold is None and arm.warm is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("ttl", "expected"), [("5m", 0.0053), ("1h", 0.0068)]) +async def test_append_only_prefix_reads_old_tokens_and_writes_extension(ttl: str, expected: float) -> None: + cache: Final = DualCache() + await _observe(cache, _body(ttl), cached_tokens=4_000) + body: Final = _body(ttl, extended=True) + arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts()) + + assert arm.cache_state == "partial" + assert arm.estimate is not None + assert arm.estimate.tokens.cache_read_input_tokens == 4_000 + assert arm.estimate.tokens.cache_creation_5m_input_tokens == (1_000 if ttl == "5m" else 0) + assert arm.estimate.tokens.cache_creation_1h_input_tokens == (1_000 if ttl == "1h" else 0) + assert arm.estimate.input_cost == pytest.approx(expected) + + +@pytest.mark.asyncio +async def test_expired_observation_estimates_a_cold_rebuild() -> None: + cache: Final = DualCache() + body: Final = _body() + await _observe(cache, body, expired=True) + arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts()) + + assert arm.cache_state == "stale" + assert arm.reason == "observation_expired" + assert arm.evidence is not None and arm.evidence.expires_at < time.time() + assert arm.estimate is not None and arm.cold is not None + assert arm.estimate.tokens.cache_read_input_tokens == 0 + assert arm.estimate.tokens.cache_creation_5m_input_tokens == 5_000 + assert arm.estimate.input_cost == arm.cold.input_cost + + +@pytest.mark.asyncio +async def test_below_model_minimum_prices_all_input_as_uncached() -> None: + body: Final = _body() + arm: Final = await endpoint.predict_arm( + _deployment(), body, _prefix(body), _CALLER, DualCache(), Counts(total=1_500, prefix=1_000) + ) + + assert arm.cache_state == "disabled" + assert arm.reason == "below_cache_minimum" + assert arm.estimate is not None + assert arm.estimate.tokens.uncached_input_tokens == 1_500 + assert arm.estimate.tokens.cache_read_input_tokens == 0 + assert arm.estimate.tokens.cache_creation_5m_input_tokens == 0 + assert arm.estimate.input_cost == pytest.approx(0.003) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("counts", [Counts(total=None), Counts(prefix=None), Counts(total=4_000)]) +async def test_unavailable_or_inconsistent_token_counts_return_null_estimates(counts: Counts) -> None: + body: Final = _body() + arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), counts) + + assert arm.cache_state == "unknown" + assert arm.reason == "token_count_unavailable" + assert arm.estimate is None and arm.cold is None and arm.warm is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("counts", [Counts(), Counts(total=1_500, prefix=1_000)]) +async def test_missing_prices_return_unknown_and_null_estimates( + monkeypatch: pytest.MonkeyPatch, counts: Counts +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "claude-cache-unpriced-5", + {"litellm_provider": "anthropic", "mode": "chat"}, + ) + body: Final = _body() + arm: Final = await endpoint.predict_arm( + _deployment("cache-prediction-unpriced", "claude-cache-unpriced-5"), + body, + _prefix(body), + _CALLER, + DualCache(), + counts, + ) + + assert arm.cache_state == "unknown" + assert arm.reason == "pricing_unavailable" + assert arm.estimate is None and arm.cold is None and arm.warm is None + + +@pytest.mark.asyncio +async def test_custom_api_base_from_environment_returns_unknown_before_counting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid") + body: Final = _body() + arm: Final = await endpoint.predict_arm( + _deployment(), body, _prefix(body), _CALLER, DualCache(), _unexpected_count + ) + + assert arm.cache_state == "unknown" + assert arm.reason == "unsupported_provider_endpoint" + assert arm.estimate is None and arm.cold is None and arm.warm is None + + +@pytest.mark.asyncio +async def test_explicit_official_api_base_overrides_custom_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid") + body: Final = _body() + arm: Final = await endpoint.predict_arm( + _deployment(api_base="https://api.anthropic.com"), body, _prefix(body), _CALLER, DualCache(), Counts() + ) + + assert arm.cache_state == "unknown" + assert arm.reason == "no_compatible_observation" + assert arm.estimate is not None + assert arm.estimate.input_cost == pytest.approx(0.0145) + + +@dataclass(frozen=True) +class _ProxyLogging: + internal_usage_cache: InternalUsageCache + parallel_limiter: CustomLogger | None + + def get_proxy_hook(self, hook: str) -> CustomLogger | None: + return self.parallel_limiter if hook == "parallel_request_limiter" else None + + +def _app( + monkeypatch: pytest.MonkeyPatch, + cache: DualCache, + *, + caller: UserAPIKeyAuth | None = None, + current_team: str | None = None, + candidate_team: str | None = None, + counts: endpoint.TokenCounter = Counts(), + limiter: CustomLogger | Literal["default"] | None = "default", +) -> FastAPI: + import litellm.proxy.proxy_server as proxy_server + + model_list: Final = [ + _deployment("opus", "claude-opus-5", team_id=current_team).model_dump(exclude_unset=True), + _deployment("sonnet", team_id=candidate_team).model_dump(exclude_unset=True), + ] + router: Final = litellm.Router(model_list=model_list) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", model_list) + monkeypatch.setattr(endpoint, "count_prompt_tokens", counts) + app: Final = FastAPI() + app.include_router(endpoint.router) + app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler) + if caller is not None: + usage_cache: Final = InternalUsageCache(cache) + configured_limiter: Final = ( + _PROXY_MaxParallelRequestsHandler_v3(usage_cache) if isinstance(limiter, str) else limiter + ) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", _ProxyLogging(usage_cache, configured_limiter)) + app.dependency_overrides[endpoint.user_api_key_auth] = lambda: caller + return app + + +async def _post( + app: FastAPI, + body: Mapping[str, JsonValue], + *, + current_deployment_id: str = "opus", + candidate_deployment_id: str = "sonnet", +) -> httpx.Response: + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + return await client.post( + "/cost/predict-cache", + json={ + "current_deployment_id": current_deployment_id, + "candidate_deployment_id": candidate_deployment_id, + "request": body, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("warm_deployment", "warm_model", "expected_delta", "expected_penalty"), + [("sonnet", "claude-sonnet-5", -0.03325, 0.0), ("opus", "claude-opus-5", 0.007, 0.0115)], +) +async def test_switch_delta_accounts_for_each_deployment_cache( + monkeypatch: pytest.MonkeyPatch, + warm_deployment: str, + warm_model: str, + expected_delta: float, + expected_penalty: float, +) -> None: + cache: Final = DualCache() + body: Final = _body() + await _observe(cache, body, deployment_id=warm_deployment, model=warm_model) + app: Final = _app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER)) + response: Final = await _post(app, body) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.switch_delta == pytest.approx(expected_delta) + assert result.cache_rebuild_penalty == pytest.approx(expected_penalty) + assert result.cache_guarantee is False + assert result.pricing_basis == "input_before_discounts_and_margins" + if warm_deployment == "sonnet": + assert result.switch.cache_state == "warm" + assert result.stay.cache_state == "unknown" + else: + assert result.stay.cache_state == "warm" + assert result.switch.cache_state == "unknown" + + +@pytest.mark.asyncio +async def test_missing_caller_identity_cannot_reuse_observations(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = DualCache() + body: Final = _body() + await _observe(cache, body) + response: Final = await _post( + _app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=None), counts=_unexpected_count), body + ) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.reason == result.switch.reason == "caller_identity_unavailable" + assert result.stay.estimate is None and result.switch.estimate is None + assert result.switch_delta is None and result.cache_rebuild_penalty is None + + +@pytest.mark.asyncio +async def test_unauthenticated_request_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "master_key", "cache-prediction-test-master-key") + response: Final = await _post(_app(monkeypatch, DualCache()), _body()) + assert response.status_code == 401, response.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arm", ["current", "candidate"]) +@pytest.mark.parametrize("caller_team", [None, "own-team"]) +@pytest.mark.parametrize("restricted", [False, True]) +async def test_foreign_and_missing_deployments_have_identical_authenticated_responses( + monkeypatch: pytest.MonkeyPatch, arm: str, caller_team: str | None, restricted: bool +) -> None: + allowed: Final = ("sonnet",) if arm == "current" else ("opus",) + app: Final = _app( + monkeypatch, + DualCache(), + caller=UserAPIKeyAuth(api_key=_CALLER, team_id=caller_team, models=list(allowed) if restricted else []), + current_team="foreign-team" if arm == "current" else None, + candidate_team="foreign-team" if arm == "candidate" else None, + counts=_unexpected_count, + ) + foreign: Final = await _post(app, _body()) + missing: Final = await _post( + app, + _body(), + current_deployment_id="missing-deployment" if arm == "current" else "opus", + candidate_deployment_id="missing-deployment" if arm == "candidate" else "sonnet", + ) + + assert foreign.status_code == missing.status_code == 404 + assert foreign.json() == missing.json() == {"detail": "Deployment not found"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deployment_team", [None, "own-team"]) +async def test_visible_public_and_own_team_deployments_remain_available( + monkeypatch: pytest.MonkeyPatch, deployment_team: str | None +) -> None: + app: Final = _app( + monkeypatch, + DualCache(), + caller=UserAPIKeyAuth(api_key=_CALLER, team_id="own-team"), + current_team=deployment_team, + candidate_team=deployment_team, + ) + response: Final = await _post(app, _body()) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.estimate is not None and result.switch.estimate is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arm", ["current", "candidate"]) +async def test_visible_deployment_outside_key_model_permissions_is_forbidden( + monkeypatch: pytest.MonkeyPatch, arm: str +) -> None: + allowed: Final = "sonnet" if arm == "current" else "opus" + denied: Final = "opus" if arm == "current" else "sonnet" + app: Final = _app(monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER, models=[allowed])) + response: Final = await _post(app, _body()) + assert response.status_code == 403, response.text + assert denied in response.text + + +@pytest.mark.asyncio +async def test_other_callers_warm_cache_is_not_prediction_evidence(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = DualCache() + body: Final = _body() + await _observe(cache, body, caller="other-caller") + response: Final = await _post(_app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER)), body) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.switch.cache_state == "unknown" + assert result.switch.reason == "no_compatible_observation" + assert result.switch.evidence is None + assert result.switch.estimate is not None + assert result.switch.estimate.tokens.cache_read_input_tokens == 0 + + +@pytest.mark.asyncio +async def test_count_failure_nulls_switch_comparison(monkeypatch: pytest.MonkeyPatch) -> None: + app: Final = _app( + monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=Counts(total=None) + ) + response: Final = await _post(app, _body()) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.reason == result.switch.reason == "token_count_unavailable" + assert result.stay.estimate is None and result.switch.estimate is None + assert result.switch_delta is None and result.cache_rebuild_penalty is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limiter", [None, CustomLogger()]) +async def test_missing_or_unsupported_limiter_returns_unknown_before_counting( + monkeypatch: pytest.MonkeyPatch, limiter: CustomLogger | None +) -> None: + app: Final = _app( + monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count, limiter=limiter + ) + response: Final = await _post(app, _body()) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.reason == result.switch.reason == "limiter_unavailable" + assert result.stay.estimate is None and result.switch.estimate is None + assert result.switch_delta is None and result.cache_rebuild_penalty is None + + +@pytest.mark.asyncio +async def test_occupied_parallel_capacity_rejects_before_provider_count(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = DualCache() + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache)) + caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1) + app: Final = _app(monkeypatch, cache, caller=caller, counts=_unexpected_count, limiter=limiter) + async with limiter.request_capacity(caller, "opus"): + response: Final = await _post(app, _body()) + + assert response.status_code == 429, response.text + assert "max_parallel_requests" in response.text + recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body()) + assert recovered.status_code == 200, recovered.text + + +@pytest.mark.asyncio +async def test_each_count_consumes_the_deployment_group_rpm_limit(monkeypatch: pytest.MonkeyPatch) -> None: + calls: Final = asyncio.Queue[str]() + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + calls.put_nowait(model) + return await Counts()(model, api_key, body) + + caller: Final = UserAPIKeyAuth(api_key=_CALLER, metadata={"model_rpm_limit": {"sonnet": 1}}) + app: Final = _app(monkeypatch, DualCache(), caller=caller, counts=count) + response: Final = await _post(app, _body()) + + assert response.status_code == 429, response.text + assert calls.qsize() == 3 + assert tuple(calls.get_nowait() for _ in range(3)) == ( + "claude-opus-5", "claude-opus-5", "claude-sonnet-5" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +async def test_each_count_preserves_auth_cached_request_tag_limits( + monkeypatch: pytest.MonkeyPatch, metadata_key: str +) -> None: + calls: Final = asyncio.Queue[str]() + caller: Final = UserAPIKeyAuth(api_key=_CALLER, metadata={"tag_rpm_limit": {"cache-cost": 1}}) + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + calls.put_nowait(model) + return await Counts()(model, api_key, body) + + async def authenticated_request(request: Request) -> UserAPIKeyAuth: + data: Final = await _read_request_body(request) + _safe_set_request_parsed_body(request, {**data, metadata_key: {"tags": ["cache-cost"]}}) + return caller + + app: Final = _app(monkeypatch, DualCache(), caller=caller, counts=count) + app.dependency_overrides[endpoint.user_api_key_auth] = authenticated_request + response: Final = await _post(app, _body()) + + assert response.status_code == 429, response.text + assert "tag_per_key" in response.text + assert calls.qsize() == 1 + assert calls.get_nowait() == "claude-opus-5" + + +@pytest.mark.asyncio +async def test_provider_counter_failure_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = DualCache() + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache)) + caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1) + + async def fail_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + raise RuntimeError("provider counter failed") + + app: Final = _app(monkeypatch, cache, caller=caller, counts=fail_count, limiter=limiter) + with pytest.raises(RuntimeError, match="provider counter failed"): + await _post(app, _body()) + recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body()) + assert recovered.status_code == 200, recovered.text + assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145) + + +@pytest.mark.asyncio +async def test_cancelled_provider_counter_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = DualCache() + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache)) + caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1) + started: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def wait_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + started.set() + await release.wait() + return await Counts()(model, api_key, body) + + app: Final = _app(monkeypatch, cache, caller=caller, counts=wait_count, limiter=limiter) + pending: Final = asyncio.create_task(_post(app, _body())) + try: + await asyncio.wait_for(started.wait(), timeout=5) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + release.set() + recovered: Final = await asyncio.wait_for(_post(app, _body()), timeout=5) + assert recovered.status_code == 200, recovered.text + assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145) + finally: + pending.cancel() + release.set() + await asyncio.gather(pending, return_exceptions=True) + + +async def _unexpected_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + pytest.fail("Unsupported prediction must return before contacting the token counter") + + +class RequestMutator(CustomLogger): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> dict[str, object]: + return {**data, "system": "Injected policy"} + + +@pytest.fixture +def request_mutator() -> Iterator[RequestMutator]: + callback: Final = RequestMutator() + litellm.logging_callback_manager.add_litellm_callback(callback) + try: + yield callback + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + + +@pytest.mark.asyncio +async def test_request_transform_callback_returns_unknown_before_token_counting( + monkeypatch: pytest.MonkeyPatch, request_mutator: RequestMutator +) -> None: + app: Final = _app( + monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count + ) + response: Final = await _post(app, _body()) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.cache_state == result.switch.cache_state == "unknown" + assert result.stay.reason == result.switch.reason == "unsupported_request_transform" + assert result.stay.estimate is None and result.switch.estimate is None + assert result.switch_delta is None and result.cache_rebuild_penalty is None + + +@pytest.mark.asyncio +async def test_key_config_returns_unknown_before_token_counting(monkeypatch: pytest.MonkeyPatch) -> None: + app: Final = _app( + monkeypatch, + DualCache(), + caller=UserAPIKeyAuth(api_key=_CALLER, config={"model_list": []}), + counts=_unexpected_count, + ) + response: Final = await _post(app, _body()) + + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.cache_state == result.switch.cache_state == "unknown" + assert result.stay.reason == result.switch.reason == "unsupported_request_transform" + assert result.stay.estimate is None and result.switch.estimate is None + assert result.switch_delta is None and result.cache_rebuild_penalty is None + + +@pytest.mark.parametrize("headers", [ + {"anthropic-version": "2099-01-01"}, + {"anthropic-beta": "future-feature"}, +]) +@pytest.mark.asyncio +async def test_unsupported_provider_headers_cannot_reuse_default_version_evidence( + monkeypatch: pytest.MonkeyPatch, headers: dict[str, str] +) -> None: + cache: Final = DualCache() + await _observe(cache, _body(), deployment_id="sonnet") + app: Final = _app( + monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response: Final = await client.post( + "/cost/predict-cache", + headers=headers, + json={"current_deployment_id": "opus", "candidate_deployment_id": "sonnet", "request": _body()}, + ) + assert response.status_code == 200, response.text + result: Final = CachePredictionResponse.model_validate(response.json()) + assert result.stay.cache_state == result.switch.cache_state == "unknown" + assert result.stay.reason == result.switch.reason == "unsupported_provider_headers" + assert result.stay.estimate is None and result.switch.estimate is None + assert result.switch_delta is None and result.cache_rebuild_penalty is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9e844b992b2..cac1005cc2f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -3387,6 +3387,35 @@ export interface paths { patch?: never; trace?: never; }; + "/cost/predict-cache": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Predict Cache Cost + * @description Compare the next native Anthropic request on two configured deployment IDs. + * + * Estimates use provider token counting and recent successful cache telemetry for this key. + * Unknown cache state uses the cold scenario when prices/counts are available. Cache observations + * do not guarantee retention. v0 supports one message-content breakpoint, text and client tools; + * system/tool-only breakpoints, thinking, images, nondefault Anthropic versions, beta headers and + * request transforms are unknown. + * Each provider count consumes one RPM unit and holds concurrency capacity; a comparison uses + * up to four counts. The legacy rate limiter returns unknown without contacting the provider. + * This endpoint does not generate tokens, prewarm caches, choose a model or alter routing. + */ + post: operations["predict_cache_cost_cost_predict_cache_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/credentials": { parameters: { query?: never; @@ -24618,6 +24647,31 @@ export interface components { /** Failed Requests */ failed_requests: number; }; + /** CacheCostScenario */ + CacheCostScenario: { + /** Input Cost */ + input_cost: number; + tokens: components["schemas"]["CacheTokenBuckets"]; + }; + /** CacheEvidence */ + CacheEvidence: { + /** + * Confidence + * @default observed + * @constant + */ + confidence: "observed"; + /** Expires At */ + expires_at: number; + /** Observed At */ + observed_at: number; + /** + * Source + * @default provider_usage + * @constant + */ + source: "provider_usage"; + }; /** CachePingResponse */ CachePingResponse: { /** Cache Type */ @@ -24635,6 +24689,59 @@ export interface components { /** Status */ status: string; }; + /** CachePredictionArm */ + CachePredictionArm: { + /** + * Cache State + * @default unknown + * @enum {string} + */ + cache_state: "warm" | "partial" | "stale" | "unknown" | "disabled"; + cold?: components["schemas"]["CacheCostScenario"] | null; + /** Deployment Id */ + deployment_id: string; + estimate?: components["schemas"]["CacheCostScenario"] | null; + evidence?: components["schemas"]["CacheEvidence"] | null; + /** Model */ + model?: string | null; + /** Reason */ + reason?: string | null; + /** Token Count Source */ + token_count_source?: "anthropic_count_tokens" | null; + warm?: components["schemas"]["CacheCostScenario"] | null; + }; + /** CachePredictionRequest */ + CachePredictionRequest: { + /** Candidate Deployment Id */ + candidate_deployment_id: string; + /** Current Deployment Id */ + current_deployment_id: string; + /** Request */ + request: { + [key: string]: components["schemas"]["JsonValue"]; + }; + }; + /** CachePredictionResponse */ + CachePredictionResponse: { + /** + * Cache Guarantee + * @default false + * @constant + */ + cache_guarantee: false; + /** Cache Rebuild Penalty */ + cache_rebuild_penalty: number | null; + /** + * Pricing Basis + * @default input_before_discounts_and_margins + * @constant + */ + pricing_basis: "input_before_discounts_and_margins"; + stay: components["schemas"]["CachePredictionArm"]; + switch: components["schemas"]["CachePredictionArm"]; + /** Switch Delta */ + switch_delta: number | null; + }; /** CacheSettingsField */ CacheSettingsField: { /** Field Default */ @@ -24716,6 +24823,29 @@ export interface components { */ status: string; }; + /** CacheTokenBuckets */ + CacheTokenBuckets: { + /** + * Cache Creation 1H Input Tokens + * @default 0 + */ + cache_creation_1h_input_tokens: number; + /** + * Cache Creation 5M Input Tokens + * @default 0 + */ + cache_creation_5m_input_tokens: number; + /** + * Cache Read Input Tokens + * @default 0 + */ + cache_read_input_tokens: number; + /** + * Uncached Input Tokens + * @default 0 + */ + uncached_input_tokens: number; + }; /** * CallTypes * @enum {string} @@ -28332,6 +28462,7 @@ export interface components { /** Updated By */ updated_by?: string | null; }; + JsonValue: unknown; /** KeyHealthResponse */ KeyHealthResponse: { /** @@ -45167,6 +45298,39 @@ export interface operations { }; }; }; + predict_cache_cost_cost_predict_cache_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["CachePredictionRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["CachePredictionResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_credentials_credentials_get: { parameters: { query?: never; From 3c2342bfd32edddadfa3429ddefa159fdd9e32c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:05:41 -0700 Subject: [PATCH 043/297] refactor(cost): return a new prompt token details wrapper when combining usage --- litellm/cost_calculator.py | 55 +++++++++++++++++++------------------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b865318f3af..50118af8a30 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2326,33 +2326,29 @@ def _combine_cached_tokens_details( ) -def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: - if not (hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details): - return - if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: - combined.prompt_tokens_details = PromptTokensDetailsWrapper() - - for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): - if ( - hasattr(usage.prompt_tokens_details, attr) - and not attr.startswith("_") - and not callable(_attribute_value(usage.prompt_tokens_details, attr)) - ): - current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 - new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 - if new_val is not None and isinstance(new_val, (int, float)): - setattr( - combined.prompt_tokens_details, - attr, - current_val + new_val, - ) - - new_cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) - if isinstance(new_cached_tokens_details, CachedTokensDetails): - combined.prompt_tokens_details.cached_tokens_details = _combine_cached_tokens_details( - getattr(combined.prompt_tokens_details, "cached_tokens_details", None), - new_cached_tokens_details, - ) +def _combine_prompt_tokens_details( + current: PromptTokensDetailsWrapper | None, new: PromptTokensDetailsWrapper +) -> PromptTokensDetailsWrapper: + base: Final = current if current is not None else PromptTokensDetailsWrapper() + base_values: Final = MappingProxyType( + {attr: getattr(base, attr) for attr in type(base).model_fields if hasattr(base, attr)} + ) + summed: Final = MappingProxyType( + { + attr: (getattr(base, attr, 0) or 0) + (getattr(new, attr) or 0) + for attr in _summable_prompt_token_fields(new) + if hasattr(new, attr) and isinstance(getattr(new, attr) or 0, (int, float)) + } + ) + new_cached_tokens_details: Final = getattr(new, "cached_tokens_details", None) + cached_tokens_details: Final = ( + _combine_cached_tokens_details(getattr(base, "cached_tokens_details", None), new_cached_tokens_details) + if isinstance(new_cached_tokens_details, CachedTokensDetails) + else getattr(base, "cached_tokens_details", None) + ) + return PromptTokensDetailsWrapper( + **MappingProxyType({**base_values, **summed, "cached_tokens_details": cached_tokens_details}) + ) class BaseTokenUsageProcessor: @@ -2381,7 +2377,10 @@ class BaseTokenUsageProcessor: and isinstance(current_val, (int, float)) ): setattr(combined, attr, current_val + new_val) - _combine_prompt_tokens_details(combined, usage) + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + combined.prompt_tokens_details = _combine_prompt_tokens_details( + getattr(combined, "prompt_tokens_details", None), usage.prompt_tokens_details + ) # Handle nested completion_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: From e4b05883627555fd72db16f0ecf376b75a6d93a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:31:53 -0700 Subject: [PATCH 044/297] fix(cost): carry cache_read_input_audio_token_cost through get_model_info Every proxy and router cost lookup goes through get_model_info, which copies cost map keys explicitly, so the new audio cache-read branch always fell back to the text cache-read rate there. Copy the key so models whose audio cache-read rate differs from the text one bill cached audio correctly. --- litellm/types/utils.py | 1 + litellm/utils.py | 1 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 17 +++++++++++++++++ tests/test_litellm/test_utils.py | 8 ++++++++ 4 files changed, 27 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9eb70e189a1..ef191d79177 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -251,6 +251,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None + cache_read_input_audio_token_cost: ReadOnly[float | None] cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing diff --git a/litellm/utils.py b/litellm/utils.py index 1a77655a5a4..d048265ecd3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5866,6 +5866,7 @@ def _get_model_info_helper( "cache_creation_input_token_cost_ultrafast", None ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), + cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( "cache_read_input_token_cost_above_200k_tokens", None diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 5d55dfb14a3..5ff1ab62698 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5235,3 +5235,20 @@ def test_cached_audio_tokens_capped_at_cached_tokens(_local_model_cost_map: None model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" ) assert prompt_cost == pytest.approx(116 * 4e-6 + (167 - 100) * 32e-6 + 100 * 4e-7) + + +def test_cached_audio_tokens_billed_at_audio_cache_rate_through_model_info_lookup(_local_model_cost_map: None) -> None: + usage = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=400, + audio_tokens=600, + cached_tokens=500, + cached_tokens_details={"text_tokens": 100, "audio_tokens": 400}, + ), + ) + + prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai") + assert prompt_cost == pytest.approx(300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 19ed31c7b22..d9ca7d9ddee 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -6430,3 +6430,11 @@ def test_completion_finishes_response_metadata_before_handing_the_response_to_th assert snapshot["litellm_call_id"] assert snapshot["response_cost"] is not None assert snapshot["api_base"] + + +def test_get_model_info_carries_cache_read_input_audio_token_cost(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + info = litellm.get_model_info("gpt-realtime-2.1-mini", custom_llm_provider="openai") + assert info["cache_read_input_audio_token_cost"] == 3e-07 + assert info["cache_read_input_token_cost"] == 6e-08 From 305caa8260fb98f821aa69eeba466a162b69c1c4 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 12 Sep 2026 23:01:41 +0000 Subject: [PATCH 045/297] test: drop unrelated reformatting from merge resolution Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_cost_calculator.py | 335 +++++++++++++++------ 1 file changed, 241 insertions(+), 94 deletions(-) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f2ee2cdbd9a..c23f5c08a70 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,3 +1,4 @@ + import json from pathlib import Path from typing import Final @@ -148,7 +149,9 @@ def test_jina_rerank_bills_total_tokens_at_input_rate_only(_local_model_cost_map def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): - _hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 1000}} + _hidden_params = { + "additional_headers": {"llm_provider-x-litellm-response-cost": 1000} + } result = response_cost_calculator( response_object=MockResponse(), @@ -204,9 +207,7 @@ def test_vertex_lyria_speech_cost( call_type=call_type, ) - expected: Final = ( - 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) - ) + expected: Final = 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) assert cost == pytest.approx(expected) @@ -333,12 +334,13 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): # Step 1: Test a model where input_cost_per_image_token is not set. # In this case the calculation should use input_cost_per_token as fallback. - assert model_info.get("input_cost_per_image_token") is None, ( - "Test case expects that input_cost_per_image_token is not set" - ) + assert ( + model_info.get("input_cost_per_image_token") is None + ), "Test case expects that input_cost_per_image_token is not set" expected_cost = ( - usage.prompt_tokens_details.audio_tokens * model_info["input_cost_per_audio_token"] + usage.prompt_tokens_details.audio_tokens + * model_info["input_cost_per_audio_token"] + usage.prompt_tokens_details.text_tokens * model_info["input_cost_per_token"] + usage.prompt_tokens_details.image_tokens * model_info["input_cost_per_token"] + usage.completion_tokens * model_info["output_cost_per_token"] @@ -373,9 +375,12 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): ) expected_cost = ( - usage.prompt_tokens_details.audio_tokens * temp_model_info_object["input_cost_per_audio_token"] - + usage.prompt_tokens_details.text_tokens * temp_model_info_object["input_cost_per_token"] - + usage.prompt_tokens_details.image_tokens * temp_model_info_object["input_cost_per_image_token"] + usage.prompt_tokens_details.audio_tokens + * temp_model_info_object["input_cost_per_audio_token"] + + usage.prompt_tokens_details.text_tokens + * temp_model_info_object["input_cost_per_token"] + + usage.prompt_tokens_details.image_tokens + * temp_model_info_object["input_cost_per_image_token"] + usage.completion_tokens * temp_model_info_object["output_cost_per_token"] ) @@ -385,11 +390,14 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): def test_transcription_cost_uses_token_pricing(_local_model_cost_map): from litellm import completion_cost + usage = Usage( prompt_tokens=14, completion_tokens=45, total_tokens=59, - prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=0, audio_tokens=14), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, audio_tokens=14 + ), ) response = TranscriptionResponse(text="demo text") response.usage = usage @@ -433,6 +441,7 @@ def test_transcription_token_pricing_is_provider_aware(_local_model_cost_map): def test_transcription_cost_falls_back_to_duration(_local_model_cost_map): from litellm import completion_cost + response = TranscriptionResponse(text="demo text") response.duration = 10.0 @@ -453,6 +462,7 @@ def test_vertex_chirp_3_transcription_cost_from_duration(_local_model_cost_map): every transcription priced to $0.00 instead of using input_cost_per_second.""" from litellm import completion_cost + response = TranscriptionResponse(text="demo text") response.duration = 18.0 @@ -476,7 +486,9 @@ def test_handle_realtime_stream_cost_calculation(): {"type": "session.created", "session": {"model": "gpt-3.5-turbo"}}, { "type": "response.done", - "response": {"usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}}, + "response": { + "usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150} + }, }, { "type": "response.done", @@ -507,7 +519,9 @@ def test_handle_realtime_stream_cost_calculation(): expected_cost = (300 * 0.0015 / 1000) + ( # input tokens (100 + 200) 150 * 0.002 / 1000 ) # output tokens (50 + 100) - assert abs(cost - expected_cost) <= 0.00075 # Allow small floating point differences + assert ( + abs(cost - expected_cost) <= 0.00075 + ) # Allow small floating point differences # Test with different model name in session results[0]["session"]["model"] = "gpt-4" @@ -587,7 +601,14 @@ def test_handle_realtime_stream_cost_calculation_stores_cost_breakdown(): assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["input_cost"] > 0 assert logging_obj.cost_breakdown["output_cost"] > 0 - assert abs(logging_obj.cost_breakdown["input_cost"] + logging_obj.cost_breakdown["output_cost"] - total_cost) < 1e-9 + assert ( + abs( + logging_obj.cost_breakdown["input_cost"] + + logging_obj.cost_breakdown["output_cost"] + - total_cost + ) + < 1e-9 + ) assert abs(logging_obj.cost_breakdown["total_cost"] - total_cost) < 1e-9 @@ -661,7 +682,9 @@ def test_realtime_logging_object_allows_null_transcript_in_conversation_item_add }, ] - usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( usage=usage, results=results, @@ -711,7 +734,9 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types(): }, ] - usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) # On unfixed code this raises pydantic ValidationError instead of returning. logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( usage=usage, @@ -723,7 +748,8 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types(): unknown_types = { r["type"] for r in logging_result.results - if r["type"] in ("rate_limits.updated", "response.function_call_arguments.delta") + if r["type"] + in ("rate_limits.updated", "response.function_call_arguments.delta") } assert unknown_types == { "rate_limits.updated", @@ -756,7 +782,9 @@ def test_realtime_transcription_duration_cost(monkeypatch): "type": "session.created", "session": { "type": "transcription", - "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}, + "audio": { + "input": {"transcription": {"model": "gpt-realtime-whisper"}} + }, }, }, { @@ -771,7 +799,9 @@ def test_realtime_transcription_duration_cost(monkeypatch): }, ] - combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) logging_obj = Logging( model="gpt-realtime-whisper", messages=[], @@ -864,7 +894,9 @@ def test_realtime_transcription_token_billed_fallback(monkeypatch): # gpt-4o-transcribe: input_cost_per_audio_token = 2.5e-06, input_cost_per_token = 2.5e-06, # output_cost_per_token = 1e-05 - model_info = litellm.get_model_info(model="gpt-4o-transcribe", custom_llm_provider="openai") + model_info = litellm.get_model_info( + model="gpt-4o-transcribe", custom_llm_provider="openai" + ) usage = { "type": "tokens", "input_tokens": 40, @@ -945,7 +977,10 @@ def test_get_transcription_model_falls_back_to_session_model(monkeypatch): mock_response=True, ) - assert result._hidden_params["response_cost"] > result_2._hidden_params["response_cost"] + assert ( + result._hidden_params["response_cost"] + > result_2._hidden_params["response_cost"] + ) model_info = router.get_deployment_model_info( model_id="my-unique-model-id", model_name="anthropic/claude-sonnet-4-5-20250929" @@ -1108,7 +1143,9 @@ def test_tiered_pricing_only_deployment_selects_router_model_id(): assert entry.get("input_cost_per_token") is None assert entry.get("tiered_pricing") is not None # The stripped shared alias must not carry tiered pricing. - assert litellm.model_cost["dashscope/qwen-tier-only-test"].get("tiered_pricing") is None + assert ( + litellm.model_cost["dashscope/qwen-tier-only-test"].get("tiered_pricing") is None + ) selected = _select_model_name_for_cost_calc( model="dashscope/qwen-tier-only-test", @@ -1188,7 +1225,9 @@ def test_azure_realtime_cost_calculator(_local_model_cost_map): combined_usage_object=Usage( prompt_tokens=100, completion_tokens=100, - prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=10, audio_tokens=90), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=10, audio_tokens=90 + ), ), custom_llm_provider="azure", litellm_model_name="my-custom-azure-deployment", @@ -1207,6 +1246,7 @@ def test_azure_audio_output_cost_calculation(_local_model_cost_map): """ from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message + # Scenario from issue #19764: # Input: 17 text tokens, 0 audio tokens # Output: 110 text tokens, 482 audio tokens @@ -1262,10 +1302,14 @@ def test_azure_audio_output_cost_calculation(_local_model_cost_map): wrong_total_cost = expected_input_cost + wrong_output_cost # Verify audio tokens are NOT charged at text rate (the bug) - assert abs(cost - wrong_total_cost) > 0.001, "Bug: Audio tokens are being charged at text token rate" + assert ( + abs(cost - wrong_total_cost) > 0.001 + ), "Bug: Audio tokens are being charged at text token rate" # Verify cost matches - assert abs(cost - expected_total_cost) < 0.0000001, f"Expected cost {expected_total_cost}, got {cost}" + assert ( + abs(cost - expected_total_cost) < 0.0000001 + ), f"Expected cost {expected_total_cost}, got {cost}" def test_default_image_cost_calculator(monkeypatch): @@ -1279,7 +1323,9 @@ def test_default_image_cost_calculator(monkeypatch): monkeypatch.setattr( litellm, "model_cost", - {"azure/bf9001cd7209f5734ecb4ab937a5a0e2ba5f119708bd68f184db362930f9dc7b": temp_object}, + { + "azure/bf9001cd7209f5734ecb4ab937a5a0e2ba5f119708bd68f184db362930f9dc7b": temp_object + }, ) args = { @@ -1495,7 +1541,9 @@ def test_gemini_25_implicit_caching_cost(): expected_cost = 0.00068708 # Allow for small floating point differences - assert abs(result - expected_cost) < 1e-8, f"Expected cost {expected_cost}, but got {result}" + assert ( + abs(result - expected_cost) < 1e-8 + ), f"Expected cost {expected_cost}, but got {result}" print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}") @@ -1566,7 +1614,9 @@ def test_log_context_cost_calculation(): # Get model info to understand the pricing from litellm import get_model_info - model_info = get_model_info(model="claude-4-sonnet-20250514", custom_llm_provider="anthropic") + model_info = get_model_info( + model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" + ) # Calculate expected cost based on actual model pricing input_cost_per_token = model_info.get("input_cost_per_token", 0) @@ -1574,8 +1624,12 @@ def test_log_context_cost_calculation(): cache_creation_cost_per_token = model_info.get("cache_creation_input_token_cost", 0) # Check if tiered pricing is applied - input_cost_above_200k = model_info.get("input_cost_per_token_above_200k_tokens", input_cost_per_token) - output_cost_above_200k = model_info.get("output_cost_per_token_above_200k_tokens", output_cost_per_token) + input_cost_above_200k = model_info.get( + "input_cost_per_token_above_200k_tokens", input_cost_per_token + ) + output_cost_above_200k = model_info.get( + "output_cost_per_token_above_200k_tokens", output_cost_per_token + ) cache_creation_above_200k = model_info.get( "cache_creation_input_token_cost_above_200k_tokens", cache_creation_cost_per_token, @@ -1583,23 +1637,31 @@ def test_log_context_cost_calculation(): print(f"DEBUG: Base input cost per token: ${input_cost_per_token:.2e}") print(f"DEBUG: Base output cost per token: ${output_cost_per_token:.2e}") - print(f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}") + print( + f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}" + ) # Handle tiered pricing - if not available, use base pricing if input_cost_above_200k is not None: - print(f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}") + print( + f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}" + ) else: print("DEBUG: No tiered input pricing available, using base pricing") input_cost_above_200k = input_cost_per_token if output_cost_above_200k is not None: - print(f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}") + print( + f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}" + ) else: print("DEBUG: No tiered output pricing available, using base pricing") output_cost_above_200k = output_cost_per_token if cache_creation_above_200k is not None: - print(f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}") + print( + f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}" + ) else: print("DEBUG: No tiered cache creation pricing available, using base pricing") cache_creation_above_200k = cache_creation_cost_per_token @@ -1613,9 +1675,13 @@ def test_log_context_cost_calculation(): print(f"DEBUG: Expected total: ${expected_total:.6f}") # Allow for small floating point differences - assert abs(result - expected_total) < 1e-6, f"Expected cost ${expected_total:.6f}, but got ${result:.6f}" + assert ( + abs(result - expected_total) < 1e-6 + ), f"Expected cost ${expected_total:.6f}, but got ${result:.6f}" - print(f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}") + print( + f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}" + ) print(f" - Input tokens (300k): ${expected_input_cost:.6f}") print(f" - Output tokens (50k): ${expected_output_cost:.6f}") print(f" - Cache creation (1k): ${expected_cache_cost:.6f}") @@ -1674,7 +1740,8 @@ def test_gemini_25_explicit_caching_cost_direct_usage(): expected_actual_cost = ( model_info["input_cost_per_token"] * usage.prompt_tokens_details.text_tokens - + model_info["cache_read_input_token_cost"] * usage.prompt_tokens_details.cached_tokens + + model_info["cache_read_input_token_cost"] + * usage.prompt_tokens_details.cached_tokens + model_info["output_cost_per_token"] * usage.completion_tokens ) @@ -1698,6 +1765,7 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import PromptTokensDetailsWrapper, Usage + # Register a custom azure_ai model with cache pricing test_model_id = "test-azure-ai-claude-model" litellm.register_model( @@ -1746,12 +1814,13 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): print(f"Output cost: {output_cost}, Expected: {expected_output_cost}") print(f"Total cost: {total_cost}") - assert abs(input_cost - expected_input_cost) < 1e-10, ( - f"Input cost mismatch: got {input_cost}, expected {expected_input_cost}" - ) - assert abs(output_cost - expected_output_cost) < 1e-10, ( - f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" - ) + assert ( + abs(input_cost - expected_input_cost) < 1e-10 + ), f"Input cost mismatch: got {input_cost}, expected {expected_input_cost}" + assert ( + abs(output_cost - expected_output_cost) < 1e-10 + ), f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" + AZURE_GPT_5_6_MAP_KEYS = ( @@ -1820,7 +1889,6 @@ def test_azure_gpt_5_6_rates_match_azure_price_page(_local_model_cost_map, model for key in token_cost_keys: assert entry[key] == pytest.approx(global_entry[key] * 1.1), key - def test_vertex_regional_deployment_costs_uplift_over_global(monkeypatch): """ Regression for https://github.com/BerriAI/litellm/issues/34393: two Vertex @@ -1903,6 +1971,7 @@ def test_cost_discount_vertex_ai(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response (use a model that exists in model_prices_and_context_window.json) response = ModelResponse( id="test-id", @@ -1931,6 +2000,7 @@ def test_cost_discount_vertex_ai(monkeypatch): custom_llm_provider="vertex_ai", ) + # Verify discount is applied (5% off means 95% of original cost) expected_cost = cost_without_discount * 0.95 assert cost_with_discount == pytest.approx(expected_cost, rel=1e-9) @@ -1948,6 +2018,7 @@ def test_cost_discount_not_applied_to_other_providers(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response for OpenAI response = ModelResponse( id="test-id", @@ -1976,6 +2047,7 @@ def test_cost_discount_not_applied_to_other_providers(monkeypatch): custom_llm_provider="openai", ) + # Costs should be the same (no discount applied to OpenAI) assert cost_with_selective_discount == cost_without_discount @@ -1991,6 +2063,7 @@ def test_cost_margin_percentage(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2019,6 +2092,7 @@ def test_cost_margin_percentage(monkeypatch): custom_llm_provider="openai", ) + # Verify margin is applied (10% margin means 110% of original cost) expected_cost = cost_without_margin * 1.10 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2036,6 +2110,7 @@ def test_cost_margin_fixed_amount(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2064,6 +2139,7 @@ def test_cost_margin_fixed_amount(monkeypatch): custom_llm_provider="openai", ) + # Verify fixed margin is applied expected_cost = cost_without_margin + 0.001 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2081,6 +2157,7 @@ def test_cost_margin_combined(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2100,7 +2177,9 @@ def test_cost_margin_combined(monkeypatch): ) # Set 8% margin + $0.0005 fixed for openai - monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"percentage": 0.08, "fixed_amount": 0.0005}}) + monkeypatch.setattr(litellm, "cost_margin_config", { + "openai": {"percentage": 0.08, "fixed_amount": 0.0005} + }) # Calculate cost with margin cost_with_margin = completion_cost( @@ -2109,6 +2188,7 @@ def test_cost_margin_combined(monkeypatch): custom_llm_provider="openai", ) + # Verify combined margin is applied expected_cost = cost_without_margin * 1.08 + 0.0005 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2126,6 +2206,7 @@ def test_cost_margin_global(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2154,6 +2235,7 @@ def test_cost_margin_global(monkeypatch): custom_llm_provider="openai", ) + # Verify global margin is applied expected_cost = cost_without_margin * 1.05 assert cost_with_global_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2171,6 +2253,7 @@ def test_cost_margin_provider_overrides_global(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2199,13 +2282,16 @@ def test_cost_margin_provider_overrides_global(monkeypatch): custom_llm_provider="openai", ) + # Verify provider-specific margin is used (not global) expected_cost = cost_without_margin * 1.10 # 10% from provider, not 5% from global assert cost_with_provider_margin == pytest.approx(expected_cost, rel=1e-9) print("✓ Cost margin provider override test passed:") print(f" - Original cost: ${cost_without_margin:.6f}") - print(f" - Cost with provider margin (10%, overrides 5% global): ${cost_with_provider_margin:.6f}") + print( + f" - Cost with provider margin (10%, overrides 5% global): ${cost_with_provider_margin:.6f}" + ) print(f" - Margin added: ${cost_with_provider_margin - cost_without_margin:.6f}") @@ -2216,6 +2302,7 @@ def test_cost_margin_with_discount(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2246,6 +2333,7 @@ def test_cost_margin_with_discount(monkeypatch): custom_llm_provider="openai", ) + # Verify: discount applied first, then margin # Base cost -> discount: base * 0.95 -> margin: (base * 0.95) * 1.10 expected_cost = base_cost * 0.95 * 1.10 @@ -2283,7 +2371,9 @@ def test_azure_image_generation_cost_calculator(): size=None, usage=ImageUsage( input_tokens=0, - input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), + input_tokens_details=ImageUsageInputTokensDetails( + image_tokens=0, text_tokens=0 + ), output_tokens=0, total_tokens=0, ), @@ -2313,6 +2403,7 @@ def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_m """Test that completion_cost extracts service_tier from completion_response object.""" from litellm import completion_cost + # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2353,18 +2444,23 @@ def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_m assert flex_cost < standard_cost, "Flex cost should be less than standard cost" flex_ratio = flex_cost / standard_cost - assert 0.45 <= flex_ratio <= 0.55, f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" + assert ( + 0.45 <= flex_ratio <= 0.55 + ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map): """Test that completion_cost extracts service_tier from usage object.""" from litellm import completion_cost + # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" # Create usage object with service_tier - usage_with_service_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + usage_with_service_tier = Usage( + prompt_tokens=1000, completion_tokens=500, total_tokens=1500 + ) # Set service_tier as an attribute on the usage object setattr(usage_with_service_tier, "service_tier", "flex") @@ -2382,7 +2478,9 @@ def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map) ) # Create usage object without service_tier - usage_without_service_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + usage_without_service_tier = Usage( + prompt_tokens=1000, completion_tokens=500, total_tokens=1500 + ) # Create ModelResponse with usage without service_tier response_standard = ModelResponse( @@ -2403,13 +2501,16 @@ def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map) assert flex_cost < standard_cost, "Flex cost should be less than standard cost" flex_ratio = flex_cost / standard_cost - assert 0.45 <= flex_ratio <= 0.55, f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" + assert ( + 0.45 <= flex_ratio <= 0.55 + ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" def test_completion_cost_service_tier_priority(_local_model_cost_map): """Test that service_tier extraction follows priority: optional_params > completion_response > usage.""" from litellm import completion_cost + # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2458,13 +2559,16 @@ def test_completion_cost_service_tier_priority(_local_model_cost_map): assert cost_from_usage > 0, "Cost from usage should be greater than 0" # Costs should be similar (all using flex) - assert abs(cost_from_params - cost_from_usage) < 1e-6, "Costs from params and usage should be similar (both flex)" + assert ( + abs(cost_from_params - cost_from_usage) < 1e-6 + ), "Costs from params and usage should be similar (both flex)" def test_completion_cost_service_tier_for_bedrock(_local_model_cost_map): """Test that Bedrock cost calculation applies service_tier-specific pricing.""" from litellm import completion_cost + model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model" litellm.register_model( model_cost={ @@ -2520,6 +2624,7 @@ def test_completion_cost_service_tier_for_anthropic(_local_model_cost_map): from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-service-tier-cost-model" litellm.register_model( model_cost={ @@ -2572,6 +2677,7 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(_local_mo from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-auto-tier-cost-model" litellm.register_model( model_cost={ @@ -2665,6 +2771,7 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier(_local_mo from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2714,6 +2821,7 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-response-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2736,7 +2844,9 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( }, reasoning_content=None, ) - response = ModelResponse(usage=usage, model=model, service_tier={"name": "priority"}) + response = ModelResponse( + usage=usage, model=model, service_tier={"name": "priority"} + ) cost = completion_cost( completion_response=response, @@ -2759,6 +2869,7 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard(_local_mo """ from litellm import completion_cost + model = "claude-test-usage-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2805,6 +2916,7 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_l ) from litellm.types.utils import PromptTokensDetailsWrapper, Usage + model = "claude-test-priority-cache-fast-model" litellm.register_model( model_cost={ @@ -2830,7 +2942,9 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_l ) usage.speed = "fast" - prompt_cost, completion_cost = anthropic_cost_per_token(model=model, usage=usage, service_tier="priority") + prompt_cost, completion_cost = anthropic_cost_per_token( + model=model, usage=usage, service_tier="priority" + ) expected_prompt = ((1000 - 200) * 6e-6 + 200 * 0.6e-6) * 2 expected_completion = 500 * 30e-6 * 2 @@ -2960,7 +3074,9 @@ def test_anthropic_fast_multiplier_only_on_models_with_fast_mode(_local_model_co "model", ["claude-sonnet-4-6", "claude-mythos-5", "claude-mythos-preview"], ) -def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models(_local_model_cost_map, monkeypatch, model): +def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models( + _local_model_cost_map, monkeypatch, model +): """ Anthropic bills every Claude 4.6+ model served with ``inference_geo="us"`` at 1.1x, and echoes that geo back in the response usage, so each of these real @@ -3025,26 +3141,28 @@ def test_gemini_cache_tokens_details_no_negative_values(): usage = VertexGeminiConfig._calculate_usage(completion_response) # Text tokens should be non-cached text only: 9402 - 9393 = 9 - assert usage.prompt_tokens_details.text_tokens == 9, ( - f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" - ) + assert ( + usage.prompt_tokens_details.text_tokens == 9 + ), f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" # Image tokens should be non-cached image only: 258 - 258 = 0 - assert usage.prompt_tokens_details.image_tokens == 0, ( - f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" - ) + assert ( + usage.prompt_tokens_details.image_tokens == 0 + ), f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" # Total cached should match - assert usage.prompt_tokens_details.cached_tokens == 9651, ( - f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" - ) + assert ( + usage.prompt_tokens_details.cached_tokens == 9651 + ), f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" # MOST IMPORTANT: text_tokens should NEVER be negative - assert usage.prompt_tokens_details.text_tokens >= 0, ( - f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - ) + assert ( + usage.prompt_tokens_details.text_tokens >= 0 + ), f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - print("✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative") + print( + "✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative" + ) def test_gemini_without_cache_tokens_details(): @@ -3112,18 +3230,18 @@ def test_gemini_implicit_caching_cost_calculation(): usage = VertexGeminiConfig._calculate_usage(completion_response) # Verify parsing - assert usage.cache_read_input_tokens == 8000, ( - f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" - ) - assert usage.prompt_tokens_details.cached_tokens == 8000, ( - f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" - ) + assert ( + usage.cache_read_input_tokens == 8000 + ), f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" + assert ( + usage.prompt_tokens_details.cached_tokens == 8000 + ), f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" # CRITICAL: text_tokens should be (10000 - 8000) = 2000, NOT 10000 # This is the fix for issue #16341 - assert usage.prompt_tokens_details.text_tokens == 2000, ( - f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" - ) + assert ( + usage.prompt_tokens_details.text_tokens == 2000 + ), f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" # Verify cost calculation uses cached token pricing response = ModelResponse( @@ -3161,7 +3279,9 @@ def test_gemini_implicit_caching_cost_calculation(): f"Cached tokens may not be using reduced pricing." ) - print("✅ Issue #16341 fix verified: Gemini implicit caching cost calculated correctly") + print( + "✅ Issue #16341 fix verified: Gemini implicit caching cost calculated correctly" + ) def test_additional_costs_only_for_azure_ai(_local_model_cost_map): @@ -3175,6 +3295,7 @@ def test_additional_costs_only_for_azure_ai(_local_model_cost_map): """ from litellm.cost_calculator import _get_additional_costs + # Non-azure_ai providers should return None result = _get_additional_costs( model="gpt-4o", @@ -3317,7 +3438,12 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_prompt_details(): }, ) - expected = (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + 100 * 0.000015 + expected = ( + (4000 - 1000 - 500) * 0.0000025 + + 1000 * 0.00000025 + + 500 * 0.000003125 + + 100 * 0.000015 + ) assert cost == pytest.approx(expected) @@ -3362,7 +3488,9 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_cache_write_tokens }, ) - expected_prompt = (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + expected_prompt = ( + (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + ) expected_completion = 100 * 0.000015 assert prompt_cost == pytest.approx(expected_prompt) @@ -3402,7 +3530,10 @@ def test_extract_cache_read_tokens_zero_when_missing(): assert _extract_cache_read_tokens({}) == 0 assert _extract_cache_read_tokens({"cache_read_input_tokens": None}) == 0 - assert _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) == 0 + assert ( + _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) + == 0 + ) def test_extract_cache_creation_tokens_anthropic_top_level(): @@ -3444,7 +3575,12 @@ def test_extract_cache_creation_tokens_zero_when_missing(): assert _extract_cache_creation_tokens({}) == 0 assert _extract_cache_creation_tokens({"cache_creation_input_tokens": None}) == 0 - assert _extract_cache_creation_tokens({"prompt_tokens_details": {"cache_write_tokens": None}}) == 0 + assert ( + _extract_cache_creation_tokens( + {"prompt_tokens_details": {"cache_write_tokens": None}} + ) + == 0 + ) def test_custom_pricing_anthropic_style_cache_tokens_not_double_counted(): @@ -3571,6 +3707,7 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message + logging_obj = Logging( model="gemini-2.5-flash", messages=[{"role": "user", "content": "Hello"}], @@ -3597,8 +3734,12 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma prompt_tokens=209, completion_tokens=3996, total_tokens=4205, - completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=3114, text_tokens=882), - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, text_tokens=109), + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=3114, text_tokens=882 + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=100, text_tokens=109 + ), ), ) @@ -3664,7 +3805,9 @@ def test_completion_cost_logs_the_rates_it_billed_at(monkeypatch): assert rates is not None assert rates.input_cost_per_token == pytest.approx(6e-6) assert rates.cache_read_input_token_cost == pytest.approx(6e-7) - assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100_000 * rates.cache_read_input_token_cost) + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx( + 100_000 * rates.cache_read_input_token_cost + ) assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(1_000 * rates.output_cost_per_token) @@ -3835,7 +3978,11 @@ def test_completion_cost_bills_interactions_api_response(): cost = completion_cost(completion_response=response, custom_llm_provider="gemini") reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] - expected = 100 * model_info["input_cost_per_token"] + 50 * model_info["output_cost_per_token"] + 25 * reasoning_rate + expected = ( + 100 * model_info["input_cost_per_token"] + + 50 * model_info["output_cost_per_token"] + + 25 * reasoning_rate + ) assert cost == pytest.approx(expected) assert cost > 0 @@ -4006,9 +4153,7 @@ def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_ assert cost == pytest.approx(3 * 4e-6 + 4014 * 4e-7 + 5 * 2e-5, rel=1e-9) -def _together_chat_response( - model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int -) -> ModelResponse: +def _together_chat_response(model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int) -> ModelResponse: return ModelResponse( id="chatcmpl-together-cache", choices=[{"finish_reason": "stop", "index": 0, "message": {"content": "acknowledged", "role": "assistant"}}], @@ -4076,8 +4221,6 @@ def test_completion_cost_together_metadata_only_model_still_uses_size_bucket(_lo ) assert cost == pytest.approx((23 + 15) * 8e-07, rel=1e-9) - - def test_select_model_name_strips_unregistered_alias_prefix(_local_model_cost_map): """A router-facing model_name alias containing "/" whose leading segment is NOT a registered provider must not be double-prefixed into a non-existent cost key. @@ -4318,7 +4461,9 @@ def test_every_one_hour_cache_write_rate_is_double_its_input_rate(): """Guard against pasting one model's 1h cache-write price onto another: every provider LiteLLM tracks (Anthropic, Bedrock, Vertex, Azure) publishes the 1h write at 2x input.""" - cost_map = json.loads((Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text()) + cost_map = json.loads( + (Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text() + ) one_hour_prefix = "cache_creation_input_token_cost_above_1hr" deviations = { (name, key): (entry["input_cost_per_token" + key[len(one_hour_prefix) :]], entry[key]) @@ -4524,7 +4669,9 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - prompt_cost, completion_cost = batch_cost_calculator(usage=usage, model="gpt-6-astra", custom_llm_provider="openai") + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, model="gpt-6-astra", custom_llm_provider="openai" + ) assert prompt_cost == pytest.approx(1000 * 5e-6) assert completion_cost == pytest.approx(500 * 2.5e-5) From 577a4e94aa383bc48271f522a2f0e00804f749a1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:09:24 -0700 Subject: [PATCH 046/297] feat(proxy): unified custom_key_policy hook for key generate, update and regenerate Adds general_settings.custom_key_policy, one coroutine that receives the operation ("generate", "update", "regenerate"), the existing key row, the effective row as it will be written, and the raw request, and can deny with a 403. It runs after the request has been normalized and before the first DB write on /key/generate, /key/service-account/generate, /key/update, /key/bulk_update, /team/key/bulk_update and /key/{key}/regenerate. The two legacy hooks keep running unchanged on the raw request. --- .../key_management_endpoints.py | 160 +++++ litellm/proxy/proxy_server.py | 9 + .../key_management_endpoints.py | 17 + .../test_key_management_endpoints.py | 621 ++++++++++++++++++ .../test_db_overlay_remote_module_scrub.py | 1 + 5 files changed, 808 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ba43943d6de..a1ce1700f68 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -148,6 +148,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyRequest, BulkUpdateKeyResponse, BulkUpdateTeamKeysRequest, + CustomKeyPolicyRequest, FailedKeyUpdate, KeySearchWhere, SuccessfulKeyUpdate, @@ -280,6 +281,7 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions: class _CustomKeyHooksModule(Protocol): user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None + user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None def _custom_key_generate_hook( @@ -294,6 +296,12 @@ def _custom_key_update_hook( return hooks.user_custom_key_update +def _custom_key_policy_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_policy + + async def _enforce_custom_key_update_policy( hook: Callable[..., Awaitable[Mapping[str, object]]] | None, data: UpdateKeyRequest, @@ -310,6 +318,113 @@ async def _enforce_custom_key_update_policy( ) +async def _enforce_custom_key_policy( + hook: Callable[..., Awaitable[Mapping[str, object]]] | None, + build_policy_request: Callable[[], CustomKeyPolicyRequest], +) -> None: + if hook is None: + return + if not inspect.iscoroutinefunction(hook): + raise ValueError("user_custom_key_policy must be a coroutine") + result: Final = await hook(build_policy_request()) + if not result.get("decision", True): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=result.get("message", "Authentication Failed - Custom Auth Rule"), + ) + + +_KEY_UPDATE_JSON_STRING_COLUMNS: Final = frozenset({"router_settings", "budget_limits"}) + +_KEY_METADATA_REQUEST_FIELDS: Final = frozenset( + (*LiteLLM_ManagementEndpoint_MetadataFields_Premium, *LiteLLM_ManagementEndpoint_MetadataFields) +) + + +def _decode_json_string_column(column: str, value: object) -> object: + if column in _KEY_UPDATE_JSON_STRING_COLUMNS and isinstance(value, str): + return json.loads(value) + return value + + +def _verification_token_from_row(row: Mapping[str, object]) -> LiteLLM_VerificationToken: + org_id: Final = row["organization_id"] if "organization_id" in row else row.get("org_id") + return LiteLLM_VerificationToken.model_validate(MappingProxyType({**row, "org_id": org_id})) + + +def _effective_key_after_update( + existing_key_row: LiteLLM_VerificationToken, + non_default_values: Mapping[str, object], +) -> LiteLLM_VerificationToken: + overlay: Final = MappingProxyType( + {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()} + ) + return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay})) + + +def _update_policy_request( + operation: Literal["update", "regenerate"], + existing_key_row: LiteLLM_VerificationToken, + non_default_values: Mapping[str, object], + request: UpdateKeyRequest | RegenerateKeyRequest, +) -> CustomKeyPolicyRequest: + return CustomKeyPolicyRequest( + operation=operation, + existing_key=_verification_token_from_row(existing_key_row.model_dump()), + effective_key=_effective_key_after_update( + existing_key_row=existing_key_row, non_default_values=non_default_values + ), + request=request, + ) + + +def _generate_budget_windows( + budget_limits: Sequence[BudgetLimitEntry] | None, +) -> tuple[Mapping[str, object], ...] | None: + if budget_limits is None: + return None + return tuple( + MappingProxyType( + { + **window.model_dump(), + "reset_at": get_budget_reset_time(budget_duration=window.budget_duration).isoformat(), + } + ) + for window in budget_limits + ) + + +def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> LiteLLM_VerificationToken: + requested: Final = data.model_dump(exclude_unset=True, exclude_none=True) + metadata_fields: Final = MappingProxyType( + {field: value for field, value in requested.items() if field in _KEY_METADATA_REQUEST_FIELDS} + ) + column_fields: Final = MappingProxyType( + {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS} + ) + request_metadata: Final = data.metadata or MappingProxyType({}) + folded_metadata: Final = {**request_metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict + columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates data_json in place + expires: Final = ( + now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None + ) + budget_reset_at: Final = ( + get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None + ) + return _verification_token_from_row( + MappingProxyType( + { + **columns, + "metadata": encrypt_callback_vars(folded_metadata), + "expires": expires, + "budget_reset_at": budget_reset_at, + "budget_limits": _generate_budget_windows(data.budget_limits), + "object_permission": None, + } + ) + ) + + def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: changed_fields: Final = MappingProxyType( { @@ -1016,6 +1131,7 @@ async def _common_key_generation_helper( litellm_changed_by: str | None, team_table: LiteLLM_TeamTableCachedObj | None, ) -> GenerateKeyResponse: + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, @@ -1164,6 +1280,16 @@ async def _common_key_generation_helper( "litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e ) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: CustomKeyPolicyRequest( + operation="generate", + existing_key=None, + effective_key=_effective_key_for_generate(data=data, now=datetime.now(timezone.utc)), + request=data, + ), + ) + # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable _budget_id = data.budget_id if prisma_client is not None and data.soft_budget is not None: @@ -2496,6 +2622,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, + user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None = None, ) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2606,6 +2733,16 @@ async def _process_single_key_update( # Prepare update data non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row) + await _enforce_custom_key_policy( + hook=user_custom_key_policy, + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=update_key_request, + ), + ) + # Update key in database if prisma_client is None: raise HTTPException( @@ -3112,6 +3249,16 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=data, + ), + ) + # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) if new_key_alias != existing_key_row.key_alias: @@ -3280,6 +3427,7 @@ async def bulk_update_keys( ) custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server) if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( @@ -3327,6 +3475,7 @@ async def bulk_update_keys( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, user_custom_key_update=custom_key_update_hook, + user_custom_key_policy=custom_key_policy_hook, ) successful_updates.append( @@ -3444,6 +3593,7 @@ async def bulk_update_team_keys( ) custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server) if prisma_client is None: raise HTTPException( @@ -3574,6 +3724,7 @@ async def bulk_update_team_keys( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, user_custom_key_update=custom_key_update_hook, + user_custom_key_policy=custom_key_policy_hook, existing_key_row=existing_by_token[db_token], ) @@ -5156,6 +5307,15 @@ async def _execute_virtual_key_regeneration( if new_key_alias != key_in_db.key_alias: _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="regenerate", + existing_key_row=key_in_db, + non_default_values=non_default_values, + request=data if data is not None else RegenerateKeyRequest(), + ), + ) update_data.update(non_default_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 606a590c24b..9d258f6fd92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -924,6 +924,7 @@ def cleanup_router_config_variables(): user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ + user_custom_key_policy, \ user_custom_sso, \ user_custom_ui_sso_sign_in_handler, \ use_background_health_checks, \ @@ -941,6 +942,7 @@ def cleanup_router_config_variables(): user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None + user_custom_key_policy = None TEAM_METADATA_VALIDATOR_REGISTRY.set(None) TEAM_METADATA_SCHEMA_REGISTRY.set(()) user_custom_sso = None @@ -2365,6 +2367,7 @@ user_custom_key_generate = None _pkce_no_redis_warning_emitted: bool = False _cp_no_redis_warning_emitted: bool = False user_custom_key_update = None +user_custom_key_policy = None user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -4250,6 +4253,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Final[dict[str, tuple[str, ...]]] = { "custom_auth", "custom_key_generate", "custom_key_update", + "custom_key_policy", "custom_team_metadata_validate", "custom_sso", "custom_ui_sso_sign_in_handler", @@ -5403,6 +5407,7 @@ class ProxyConfig: user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ + user_custom_key_policy, \ user_custom_sso, \ user_custom_ui_sso_sign_in_handler, \ use_background_health_checks, \ @@ -5940,6 +5945,10 @@ class ProxyConfig: if custom_key_update is not None: user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) + custom_key_policy: Final = general_settings.get("custom_key_policy", None) + if custom_key_policy is not None: + user_custom_key_policy = get_instance_fn(value=custom_key_policy, config_file_path=config_file_path) + custom_team_metadata_validate: Final = general_settings.get("custom_team_metadata_validate", None) TEAM_METADATA_VALIDATOR_REGISTRY.set( get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path) diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 9fb5bea81e3..35abb8fb718 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -4,6 +4,9 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict +from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy._types import GenerateKeyRequest, RegenerateKeyRequest, UpdateKeyRequest +from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains @@ -123,3 +126,17 @@ class BulkUpdateTeamKeysRequest(BaseModel): if not has_key_ids and not self.all_keys_in_team: raise ValueError("Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`.") return self + + +CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"] + + +class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase): + """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request.""" + + model_config = ConfigDict(protected_namespaces=(), frozen=True) + + operation: CustomKeyPolicyOperation + existing_key: LiteLLM_VerificationToken | None + effective_key: LiteLLM_VerificationToken + request: GenerateKeyRequest | UpdateKeyRequest | RegenerateKeyRequest diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 000a9dc2f1a..2b8c62523b9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1,3 +1,4 @@ +from contextlib import ExitStack from typing import Final import json from datetime import datetime, timedelta, timezone @@ -22,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LiteLLM_VerificationToken, + LiteLLMKeyType, LitellmUserRoles, Member, ProxyException, @@ -38,6 +40,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_project_key_limits, _check_team_key_limits, _common_key_generation_helper, + _effective_key_after_update, + _effective_key_for_generate, + _enforce_custom_key_policy, _enforce_upperbound_key_params, _execute_virtual_key_regeneration, _get_and_validate_existing_key, @@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_key_team_change, ) from litellm.proxy.proxy_server import app +from litellm.types.proxy.management_endpoints.key_management_endpoints import CustomKeyPolicyRequest client = TestClient(app) @@ -12178,6 +12184,586 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +_POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" +_POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" +_POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} + + +def _seven_day_policy(received: list[CustomKeyPolicyRequest]): + async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + expires = policy_request.effective_key.expires + if isinstance(expires, datetime) and expires > datetime.now(timezone.utc) + timedelta(days=7): + return {"decision": False, "message": _POLICY_DENIAL_MESSAGE} + return {"decision": True} + + return policy + + +def _assert_expires_in(effective_key: LiteLLM_VerificationToken, duration: str) -> None: + expires = effective_key.expires + assert isinstance(expires, datetime) + assert expires.tzinfo is not None + expected = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration=duration)) + assert abs((expires - expected).total_seconds()) < 60 + + +def _regenerate_policy_mocks(policy, insert_deprecated_key: AsyncMock, persist: AsyncMock) -> ExitStack: + stack = ExitStack() + stack.enter_context( + patch( # test-quality-ok: deterministic token setup for the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ) + ) + stack.enter_context( + patch( # test-quality-ok: grace-period write must not run on a denied regenerate + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + insert_deprecated_key, + ) + ) + stack.enter_context( + patch( # test-quality-ok: archival write must not run on a denied regenerate + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + persist, + ) + ) + stack.enter_context( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) + ) + stack.enter_context( + patch( # test-quality-ok: rotation callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook + ) + return stack + + +async def _regenerate_under_policy(mock_prisma_client, existing_key, data): + return await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_regenerate_rejects_when_custom_key_policy_denies_the_effective_expiry(): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + insert_deprecated_key = AsyncMock() + persist = AsyncMock() + + with _regenerate_policy_mocks(_seven_day_policy(received), insert_deprecated_key, persist): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="3000d")) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + insert_deprecated_key.assert_not_awaited() + persist.assert_not_awaited() + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert [policy_request.operation for policy_request in received] == ["regenerate"] + assert received[0].existing_key is not None + assert received[0].existing_key.token == "abc123" + assert isinstance(received[0].request, RegenerateKeyRequest) + assert received[0].request.duration == "3000d" + _assert_expires_in(received[0].effective_key, "3000d") + + +@pytest.mark.asyncio +async def test_regenerate_within_custom_key_policy_rotates_the_key(): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + persist = AsyncMock() + + with _regenerate_policy_mocks(_seven_day_policy(received), AsyncMock(), persist): + await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="5d")) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + persist.assert_awaited_once() + assert persist.call_args.kwargs["keys"] == [existing_key] + assert [policy_request.operation for policy_request in received] == ["regenerate"] + _assert_expires_in(received[0].effective_key, "5d") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +async def test_regenerate_without_changes_still_runs_custom_key_policy(data): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + + async def freeze_rotation(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + return {"decision": False, "message": "key rotation is frozen"} + + with _regenerate_policy_mocks(freeze_rotation, AsyncMock(), AsyncMock()): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, existing_key, data) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "key rotation is frozen" + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert [policy_request.operation for policy_request in received] == ["regenerate"] + assert received[0].existing_key == existing_key + assert received[0].effective_key == existing_key + + +def _policy_existing_team_key() -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token=_POLICY_HASHED_TOKEN, user_id="test-user", team_id="team-a", max_budget=200.0 + ) + + +def _setup_update_key_fn_policy_mocks(monkeypatch, existing_key: LiteLLM_VerificationToken) -> AsyncMock: + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"max_budget": 50.0, "team_id": "team-a"}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", AsyncMock(return_value=None) + ) + return mock_prisma_client + + +def _assert_update_policy_request(policy_request: CustomKeyPolicyRequest, request: UpdateKeyRequest) -> None: + assert policy_request.operation == "update" + assert policy_request.request is request + assert policy_request.existing_key is not None + assert policy_request.existing_key.max_budget == 200.0 + assert policy_request.effective_key.team_id == "team-a" + assert policy_request.effective_key.user_id == "test-user" + assert policy_request.effective_key.max_budget == 50.0 + _assert_expires_in(policy_request.effective_key, request.duration or "") + + +@pytest.mark.asyncio +async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + received: list[CustomKeyPolicyRequest] = [] + policy = _seven_day_policy(received) + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy), # test-quality-ok: inject policy hook + ): + await update_key_fn( + request=MagicMock(), + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + mock_prisma_client.update_data.assert_awaited_once() + assert len(received) == 1 + _assert_update_policy_request(received[0], data) + + +@pytest.mark.asyncio +async def test_update_key_fn_rejects_when_custom_key_policy_denies(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + received: list[CustomKeyPolicyRequest] = [] + policy = _seven_day_policy(received) + + with patch("litellm.proxy.proxy_server.user_custom_key_policy", policy): # test-quality-ok: inject policy hook + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert str(exc_info.value.code) == "403" + assert exc_info.value.message == _POLICY_DENIAL_MESSAGE + mock_prisma_client.update_data.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["update"] + _assert_expires_in(received[0].effective_key, "3000d") + + +async def _process_single_key_update_under_policy(prisma_client: AsyncMock, data: UpdateKeyRequest, policy): + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: update callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + return await _process_single_key_update( + update_key_request=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + prisma_client=prisma_client, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + llm_router=None, + existing_key_row=_policy_existing_team_key(), + user_custom_key_policy=policy, + ) + + +@pytest.mark.asyncio +async def test_process_single_key_update_runs_custom_key_policy_on_the_effective_row(): + mock_prisma_client = AsyncMock() + updated_row = MagicMock() + updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + received: list[CustomKeyPolicyRequest] = [] + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + + result = await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received)) + + assert result["max_budget"] == 50.0 + mock_prisma_client.update_data.assert_awaited_once() + assert len(received) == 1 + _assert_update_policy_request(received[0], data) + + +@pytest.mark.asyncio +async def test_process_single_key_update_rejects_when_custom_key_policy_denies(): + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock() + received: list[CustomKeyPolicyRequest] = [] + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received)) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + mock_prisma_client.update_data.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["update"] + + +@pytest.mark.asyncio +async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + ) + + existing_keys = [ + LiteLLM_VerificationToken(token="test-key-1", user_id="user-123", max_budget=None), + LiteLLM_VerificationToken(token="test-key-2", user_id="user-123", max_budget=50.0), + ] + updated_row = MagicMock() + updated_row.model_dump.return_value = {"user_id": "user-123", "max_budget": 100.0} + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=existing_keys) + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + mock_prisma_client.get_data = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + received: list[CustomKeyPolicyRequest] = [] + + async def cap_max_budget(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + max_budget = policy_request.effective_key.max_budget + if max_budget is not None and max_budget > 100: + return {"decision": False, "message": "max_budget must be 100 or less"} + return {"decision": True} + + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", cap_max_budget) + + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: update callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + response = await bulk_update_keys( + data=BulkUpdateKeyRequest( + keys=[ + BulkUpdateKeyRequestItem(key="test-key-1", max_budget=100.0), + BulkUpdateKeyRequestItem(key="test-key-2", max_budget=500.0), + ] + ), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert [update.key for update in response.successful_updates] == ["test-key-1"] + assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [ + ("test-key-2", "max_budget must be 100 or less") + ] + assert mock_prisma_client.update_data.await_count == 1 + assert [policy_request.operation for policy_request in received] == ["update", "update"] + assert [policy_request.effective_key.max_budget for policy_request in received] == [100.0, 500.0] + assert [ + policy_request.existing_key.max_budget if policy_request.existing_key is not None else "missing" + for policy_request in received + ] == [None, 50.0] + + +def _policy_generate_prisma() -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1")) + mock_prisma.jsonify_object = MagicMock(side_effect=lambda data: json.loads(data) if isinstance(data, str) else data) + return mock_prisma + + +def _generate_policy_mocks(mock_prisma: MagicMock, generate_key_helper: AsyncMock, policy) -> ExitStack: + stack = ExitStack() + stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)) # test-quality-ok: fake DB + stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", None)) # test-quality-ok: no router in test + stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) # test-quality-ok: premium fields + stack.enter_context(patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")) # test-quality-ok: admin + stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())) # test-quality-ok: cache + stack.enter_context( + patch( # test-quality-ok: the key write must not run on a denied generate + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + generate_key_helper, + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook + ) + return stack + + +def _generate_request(duration: str, organization_id: str | None) -> GenerateKeyRequest: + return GenerateKeyRequest( + duration=duration, + organization_id=organization_id, + guardrails=["g1"], + tags=["t1"], + soft_budget=10.0, + max_budget=20.0, + ) + + +def _assert_generate_policy_request( + policy_request: CustomKeyPolicyRequest, duration: str, organization_id: str | None +) -> None: + assert policy_request.operation == "generate" + assert policy_request.existing_key is None + assert policy_request.effective_key.org_id == organization_id + assert policy_request.effective_key.max_budget == 20.0 + assert policy_request.effective_key.metadata["guardrails"] == ["g1"] + assert policy_request.effective_key.metadata["tags"] == ["t1"] + _assert_expires_in(policy_request.effective_key, duration) + + +@pytest.mark.asyncio +async def test_generate_key_rejects_when_custom_key_policy_denies_before_any_write(): + mock_prisma = _policy_generate_prisma() + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + data = _generate_request("3000d", organization_id="org-1") + + with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)): + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None + ) + + assert str(exc_info.value.code) == "403" + assert exc_info.value.message == _POLICY_DENIAL_MESSAGE + mock_prisma.db.litellm_budgettable.create.assert_not_awaited() + generate_key_helper.assert_not_awaited() + assert len(received) == 1 + _assert_generate_policy_request(received[0], "3000d", organization_id="org-1") + assert received[0].request is data + assert data.duration == "3000d" + assert data.guardrails == ["g1"] + assert data.tags == ["t1"] + assert data.organization_id == "org-1" + + +@pytest.mark.asyncio +async def test_generate_key_within_custom_key_policy_creates_the_key(): + mock_prisma = _policy_generate_prisma() + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + data = _generate_request("5d", organization_id=None) + + with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)): + await generate_key_fn( + data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None + ) + + mock_prisma.db.litellm_budgettable.create.assert_awaited_once() + generate_key_helper.assert_awaited_once() + assert len(received) == 1 + _assert_generate_policy_request(received[0], "5d", organization_id=None) + assert received[0].request is data + + +@pytest.mark.asyncio +async def test_service_account_generate_rejects_when_custom_key_policy_denies(): + from litellm.proxy.management_endpoints.key_management_endpoints import generate_service_account_key_fn + + mock_prisma = _policy_generate_prisma() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock()) + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + + with ( + _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)), + patch( # test-quality-ok: team lookup is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id="team-1", duration="3000d"), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + generate_key_helper.assert_not_awaited() + mock_prisma.db.litellm_budgettable.create.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["generate"] + assert received[0].existing_key is None + assert received[0].effective_key.team_id == "team-1" + assert received[0].effective_key.user_id is None + _assert_expires_in(received[0].effective_key, "3000d") + + +@pytest.mark.asyncio +async def test_effective_key_after_update_decodes_json_string_columns_and_keeps_omitted_fields(): + existing_key = LiteLLM_VerificationToken(token="tok", user_id="u1", team_id="team-a") + non_default_values = await prepare_key_update_data( + data=UpdateKeyRequest( + key="tok", router_settings={"num_retries": 3}, budget_limits=[{"budget_duration": "1d", "max_budget": 2.0}] + ), + existing_key_row=existing_key, + ) + assert isinstance(non_default_values["router_settings"], str) + assert isinstance(non_default_values["budget_limits"], str) + + effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values) + + assert effective_key.router_settings == {"num_retries": 3} + assert effective_key.budget_limits is not None + assert effective_key.budget_limits[0]["max_budget"] == 2.0 + assert effective_key.budget_limits[0]["budget_duration"] == "1d" + assert effective_key.budget_limits[0]["reset_at"] is not None + assert effective_key.team_id == "team-a" + assert effective_key.user_id == "u1" + + +@pytest.mark.asyncio +async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration(): + existing_key = LiteLLM_VerificationToken(token="tok", expires=datetime(2027, 1, 1, tzinfo=timezone.utc)) + non_default_values = await prepare_key_update_data( + data=UpdateKeyRequest(key="tok", duration="-1"), existing_key_row=existing_key + ) + + effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values) + + assert effective_key.expires is None + + +def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it(): + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + data = GenerateKeyRequest( + duration="5d", + organization_id="org-1", + metadata={"a": 1}, + guardrails=["g1"], + tags=["t1"], + budget_duration="1d", + max_budget=3.0, + key_type=LiteLLMKeyType.LLM_API, + ) + + effective_key = _effective_key_for_generate(data=data, now=now) + + assert effective_key.expires == now + timedelta(days=5) + assert effective_key.org_id == "org-1" + assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]} + assert effective_key.max_budget == 3.0 + assert effective_key.budget_duration == "1d" + assert effective_key.budget_reset_at is not None + assert effective_key.key_type == "llm_api" + assert effective_key.allowed_routes == ["llm_api_routes"] + assert data.metadata == {"a": 1} + assert data.guardrails == ["g1"] + assert data.tags == ["t1"] + assert data.duration == "5d" + + +def test_effective_key_for_generate_without_duration_never_expires(): + effective_key = _effective_key_for_generate( + data=GenerateKeyRequest(), now=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + assert effective_key.expires is None + assert effective_key.budget_reset_at is None + assert effective_key.key_type == "default" + + +def _policy_request_for_generate() -> CustomKeyPolicyRequest: + return CustomKeyPolicyRequest( + operation="generate", + existing_key=None, + effective_key=LiteLLM_VerificationToken(token="tok"), + request=GenerateKeyRequest(), + ) + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_rejects_a_sync_hook(): + def sync_hook(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + return {"decision": True} + + with pytest.raises(ValueError, match="user_custom_key_policy must be a coroutine"): + await _enforce_custom_key_policy(hook=sync_hook, build_policy_request=_policy_request_for_generate) + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_uses_the_default_denial_message(): + async def deny(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + return {"decision": False} + + with pytest.raises(HTTPException) as exc_info: + await _enforce_custom_key_policy(hook=deny, build_policy_request=_policy_request_for_generate) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule" + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ @@ -18348,3 +18934,38 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): ) assert exc.value.status_code == 403 assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_runs_custom_key_policy_per_key(monkeypatch): + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + keys = [_make_team_key("tok-a"), _make_team_key("tok-b")] + mock = _setup_team_keys_mocks( + monkeypatch, find_many=keys, update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}) + ) + received: list[CustomKeyPolicyRequest] = [] + + async def freeze_tok_b(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + if policy_request.existing_key is not None and policy_request.existing_key.token == "tok-b": + return {"decision": False, "message": "tok-b is frozen"} + return {"decision": True} + + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", freeze_tok_b) + + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", key_ids=["tok-a", "tok-b"], update_fields=KeyUpdateFields(max_budget=50.0) + ) + ) + + assert [update.key for update in response.successful_updates] == ["tok-a"] + assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [("tok-b", "tok-b is frozen")] + mock.update_data.assert_awaited_once() + assert [policy_request.operation for policy_request in received] == ["update", "update"] + assert [policy_request.effective_key.max_budget for policy_request in received] == [50.0, 50.0] + assert [policy_request.effective_key.team_id for policy_request in received] == ["team-abc", "team-abc"] diff --git a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py index 100ba653f3a..0072997a0d9 100644 --- a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py +++ b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py @@ -43,6 +43,7 @@ def test_litellm_settings_callback_list_strips_remote_urls(field): "custom_auth", "custom_key_generate", "custom_key_update", + "custom_key_policy", "custom_sso", "custom_ui_sso_sign_in_handler", ], From 83594427fca6ca448426a7fd219ce68a832ba26c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:10:10 -0700 Subject: [PATCH 047/297] fix(vertex-live): price each grounded turn's query fee on the /vertex_ai/live passthrough --- ...tex_ai_live_passthrough_logging_handler.py | 130 +++++++++++++++--- .../test_vertex_ai_live_passthrough.py | 63 ++++++++- 2 files changed, 174 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0e2eb60704d..0ac04654e30 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -7,11 +7,12 @@ Supports different modalities: text, audio, video, and web search. from collections.abc import Mapping, Sequence from datetime import datetime -from itertools import chain +from itertools import chain, pairwise from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( BasePassthroughLoggingHandler, ) @@ -20,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrou ) from litellm.types.utils import ( CompletionTokensDetailsWrapper, + CostBreakdown, LlmProviders, ModelResponse, PromptTokensDetailsWrapper, @@ -60,6 +62,35 @@ def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[s ) +def _turns(websocket_messages: Sequence[object]) -> tuple[tuple[object, ...], ...]: + """Split a session at every ``usageMetadata`` frame; frames after the last one never got their usage.""" + closes: Final = tuple( + index + 1 + for index, message in enumerate(websocket_messages) + if isinstance(message, Mapping) and isinstance(message.get("usageMetadata"), dict) + ) + return tuple(tuple(websocket_messages[start:end]) for start, end in pairwise((0, *closes))) + + +_SummedField: TypeAlias = Literal[ + "input_cost", + "output_cost", + "tool_usage_cost", + "cache_read_cost", + "cache_creation_cost", + "reasoning_cost", + "original_cost", + "discount_amount", + "margin_fixed_amount", + "margin_total_amount", +] + + +def _summed(breakdowns: Sequence[CostBreakdown], field: _SummedField) -> float | None: + values: Final = tuple(value for breakdown in breakdowns if (value := breakdown.get(field)) is not None) + return sum(values) if values else None + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -141,7 +172,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _extract_usage_metadata_from_websocket_messages( - websocket_messages: list[dict], + websocket_messages: Sequence[object], ) -> dict | None: """ Extract and aggregate usage metadata from a list of WebSocket messages. @@ -158,9 +189,11 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Dictionary containing aggregated usage metadata, or None if not found """ snapshots: Final = tuple( - message["usageMetadata"] + metadata for message in websocket_messages - if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict) + if isinstance(message, Mapping) + for metadata in (message.get("usageMetadata"),) + if isinstance(metadata, dict) ) if not snapshots: @@ -247,10 +280,72 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): ) return usage + def _session_usage(self, websocket_messages: Sequence[object], model: str) -> Usage | None: + usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) + if usage_metadata is None: + return None + return self._create_usage_object_from_metadata( + usage_metadata=usage_metadata, + grounding_metadata=_grounding_metadata(websocket_messages), + model=model, + ) + + def _turn_cost( + self, + turn: Sequence[object], + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> tuple[float, CostBreakdown] | None: + usage: Final = self._session_usage(turn, model) + if usage is None: + return None + cost: Final = logging_obj._response_cost_calculator( # pyright: ignore[reportPrivateUsage] # the call's own calculator keeps custom pricing and the deployment's region in step with the spend row + result=ModelResponse(model=model, usage=usage), + litellm_model_name=model, + ) + if cost is None: + return None + breakdown: Final = logging_obj.cost_breakdown + return None if breakdown is None else (cost, breakdown) + + def _session_cost( + self, + websocket_messages: Sequence[object], + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> float | None: + """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.""" + turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages)) + priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None) + if not priced or len(priced) != len(turn_costs): + return None + breakdowns: Final = tuple(breakdown for _, breakdown in priced) + first: Final = breakdowns[0] + total_cost: Final = sum(cost for cost, _ in priced) + logging_obj.set_cost_breakdown( + input_cost=_summed(breakdowns, "input_cost") or 0.0, + output_cost=_summed(breakdowns, "output_cost") or 0.0, + total_cost=total_cost, + cost_for_built_in_tools_cost_usd_dollar=_summed(breakdowns, "tool_usage_cost") or 0.0, + original_cost=_summed(breakdowns, "original_cost"), + discount_percent=first.get("discount_percent"), + discount_amount=_summed(breakdowns, "discount_amount"), + margin_percent=first.get("margin_percent"), + margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"), + margin_total_amount=_summed(breakdowns, "margin_total_amount"), + cache_read_cost=_summed(breakdowns, "cache_read_cost"), + cache_creation_cost=_summed(breakdowns, "cache_creation_cost"), + reasoning_cost=_summed(breakdowns, "reasoning_cost"), + service_tier=first.get("service_tier"), + data_residency=first.get("data_residency"), + vertex_location=first.get("vertex_location"), + ) + return total_cost + def vertex_ai_live_passthrough_handler( self, - websocket_messages: list[dict], - logging_obj, + websocket_messages: Sequence[object], + logging_obj: LiteLLMLoggingObj, url_route: str, start_time: datetime, end_time: datetime, @@ -274,28 +369,25 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ try: # Extract model from request body or kwargs - model: Final = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09") + requested_model: Final = kwargs.get("model") + model: Final = ( + requested_model if isinstance(requested_model, str) else "gemini-2.0-flash-live-preview-04-09" + ) custom_llm_provider: Final = kwargs.get("custom_llm_provider", "vertex_ai") verbose_proxy_logger.debug( "Vertex AI Live API model: %s, custom_llm_provider: %s", model, custom_llm_provider ) - # Extract usage metadata from WebSocket messages - usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) + usage: Final = self._session_usage(websocket_messages, model) - if not usage_metadata: + if usage is None: verbose_proxy_logger.warning("No usage metadata found in Vertex AI Live API WebSocket messages") return { "result": None, "kwargs": kwargs, } - # Create Usage object for standard LiteLLM logging - usage: Final = self._create_usage_object_from_metadata( - usage_metadata=usage_metadata, - grounding_metadata=_grounding_metadata(websocket_messages), - model=model, - ) + response_cost: Final = self._session_cost(websocket_messages, model, logging_obj) # Create a mock ModelResponse for standard logging litellm_model_response: Final = ModelResponse( @@ -306,6 +398,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): usage=usage, choices=[], ) + if response_cost is not None: + litellm_model_response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads the cost off the response's hidden params; the constructor's hidden_params kwarg is reset by pydantic kwargs["model"] = model kwargs["custom_llm_provider"] = custom_llm_provider @@ -314,7 +408,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): import re allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$") - safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" + safe_model: Final = model if allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( "Vertex AI Live API passthrough cost tracking - Model: %s, " "Prompt tokens: %s %s, Completion tokens: %s %s", diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 1815ff134aa..dbb6c4cd702 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -24,7 +24,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.utils import LlmProviders, Usage +from litellm.types.utils import CostBreakdown, LlmProviders, Usage from litellm.proxy._types import UserAPIKeyAuth @@ -47,6 +47,7 @@ class TestVertexAILivePassthroughLoggingHandler: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock @pytest.fixture @@ -474,6 +475,64 @@ class TestVertexAILivePassthroughLoggingHandler: assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded" + def _priced_logging_obj(self) -> LiteLLMLoggingObj: + """A real logging object, since the session's price is handed to it turn by turn.""" + logging_obj = LiteLLMLoggingObj( + model=self.NATIVE_AUDIO_MODEL, + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="live-session", + function_id="live", + ) + logging_obj.update_environment_variables( + model=self.NATIVE_AUDIO_MODEL, + user="u", + optional_params={}, + litellm_params={}, + call_type="pass_through_endpoint", + ) + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + return logging_obj + + def _billed_session( + self, handler: VertexAILivePassthroughLoggingHandler, messages: list[dict[str, object]] + ) -> tuple[float, CostBreakdown]: + logging_obj = self._priced_logging_obj() + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=self.NATIVE_AUDIO_MODEL, + custom_llm_provider="vertex_ai", + ) + assert result["result"] is not None, "the handler must produce a usage-bearing response to bill" + assert logging_obj.cost_breakdown is not None, "the session's price must reach the logging object" + return result["result"]._hidden_params["response_cost"], logging_obj.cost_breakdown + + def test_each_grounded_turn_pays_its_own_query_fee(self, handler): + """Google charges the grounding fee per grounded prompt, not per session. + + Summing the session into one usage collapsed two grounded turns into one query, so the + second question was answered for free. The bill now grows by one fee per grounded turn. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + + plain_cost, _ = self._billed_session(handler, [head, turn, turn]) + one_cost, one_breakdown = self._billed_session(handler, [head, grounding, turn, turn]) + two_cost, two_breakdown = self._billed_session(handler, [head, grounding, turn, grounding, turn]) + + fee = one_cost - plain_cost + assert fee > 0, "a grounded turn must cost more than the same tokens ungrounded" + assert two_cost - plain_cost == pytest.approx(2 * fee), "two grounded turns must pay the fee twice" + assert two_breakdown["total_cost"] == pytest.approx(two_cost) + assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. @@ -676,6 +735,7 @@ class TestVertexAILivePassthroughIntegration: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock @patch( @@ -809,6 +869,7 @@ class TestVertexAILivePassthroughErrorHandling: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock def test_invalid_websocket_messages_format(self): From 710d4ae2a3aea57f78decbe856dca64e4b7e8224 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:33:31 -0700 Subject: [PATCH 048/297] fix(vertex-live): charge the fixed cost margin once per Live session --- ...tex_ai_live_passthrough_logging_handler.py | 18 +++++++++++---- .../test_vertex_ai_live_passthrough.py | 23 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0ac04654e30..57ec960b032 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -314,14 +314,24 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): model: str, logging_obj: LiteLLMLoggingObj, ) -> float | None: - """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.""" + """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice. + + The fixed cost margin is a flat per-request fee, so the session's single spend row carries it once + rather than once per turn. + """ turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages)) priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None) if not priced or len(priced) != len(turn_costs): return None breakdowns: Final = tuple(breakdown for _, breakdown in priced) first: Final = breakdowns[0] - total_cost: Final = sum(cost for cost, _ in priced) + fixed_margin: Final = first.get("margin_fixed_amount") or 0.0 + duplicated_fixed_margin: Final = fixed_margin * (len(priced) - 1) + total_cost: Final = sum(cost for cost, _ in priced) - duplicated_fixed_margin + summed_margin_total: Final = _summed(breakdowns, "margin_total_amount") + margin_total_amount: Final = ( + None if summed_margin_total is None else summed_margin_total - duplicated_fixed_margin + ) logging_obj.set_cost_breakdown( input_cost=_summed(breakdowns, "input_cost") or 0.0, output_cost=_summed(breakdowns, "output_cost") or 0.0, @@ -331,8 +341,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): discount_percent=first.get("discount_percent"), discount_amount=_summed(breakdowns, "discount_amount"), margin_percent=first.get("margin_percent"), - margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"), - margin_total_amount=_summed(breakdowns, "margin_total_amount"), + margin_fixed_amount=first.get("margin_fixed_amount"), + margin_total_amount=margin_total_amount, cache_read_cost=_summed(breakdowns, "cache_read_cost"), cache_creation_cost=_summed(breakdowns, "cache_creation_cost"), reasoning_cost=_summed(breakdowns, "reasoning_cost"), diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index dbb6c4cd702..70c2fda369b 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -13,6 +13,7 @@ from typing import Dict, List, Any, Optional import pytest import httpx +import litellm from typing_extensions import NotRequired, ReadOnly, TypedDict # Add the parent directory to the system path @@ -533,6 +534,28 @@ class TestVertexAILivePassthroughLoggingHandler: assert two_breakdown["total_cost"] == pytest.approx(two_cost) assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_the_fixed_cost_margin_is_charged_once_per_session(self, handler): + """A fixed cost margin is a flat per-request fee, and a Live session is one spend row. + + Pricing each turn on its own applied the fixed margin per turn, so a two-turn session paid it + twice. The session now carries the fixed margin once no matter how many turns it billed. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + messages = [head, grounding, turn, grounding, turn] + + plain_cost, _ = self._billed_session(handler, messages) + + fixed_amount = 0.01 + with patch.object(litellm, "cost_margin_config", {"vertex_ai": {"fixed_amount": fixed_amount}}): + margined_cost, breakdown = self._billed_session(handler, messages) + + assert margined_cost - plain_cost == pytest.approx( + fixed_amount + ), "a two-turn session must add the fixed margin once, not once per billed turn" + assert breakdown["margin_fixed_amount"] == pytest.approx(fixed_amount) + assert breakdown["margin_total_amount"] == pytest.approx(fixed_amount) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. From 59c4cf94397207fa0604d58dac31027b94b8467f Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 12 Sep 2026 23:37:07 +0000 Subject: [PATCH 049/297] refactor(responses): build InputTokensDetails without post-construction mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index c75e5d0f5ea..14957a5e5aa 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2748,17 +2748,20 @@ class LiteLLMCompletionResponsesConfig: cache_write_tokens: Final = getattr(prompt_details, "cache_write_tokens", None) or getattr( prompt_details, "cache_creation_tokens", None ) - input_tokens_details: Final = InputTokensDetails( + cache_write_extra: Final[Mapping[str, int]] = ( + MappingProxyType({"cache_write_tokens": cache_write_tokens}) + if cache_write_tokens is not None + else MappingProxyType({}) + ) + response_usage.input_tokens_details = InputTokensDetails( cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, text_tokens=prompt_details.text_tokens, audio_tokens=prompt_details.audio_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), + **cache_write_extra, ) - if cache_write_tokens is not None: - setattr(input_tokens_details, "cache_write_tokens", cache_write_tokens) - response_usage.input_tokens_details = input_tokens_details # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: From fc62df33c2a0b049e8b48baad9bf97cd87aa57d3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:42:20 -0700 Subject: [PATCH 050/297] fix(proxy): resolve rotation and permission fields before the key policy and pin the effective-row contract --- .../key_management_endpoints.py | 38 +++++++---- .../key_management_endpoints.py | 13 +++- .../test_key_management_endpoints.py | 68 ++++++++++++++++++- 3 files changed, 100 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a1ce1700f68..841c22da050 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -359,7 +359,9 @@ def _effective_key_after_update( overlay: Final = MappingProxyType( {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()} ) - return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay})) + return _verification_token_from_row( + MappingProxyType({**existing_key_row.model_dump(), **overlay, "object_permission": None}) + ) def _update_policy_request( @@ -381,7 +383,7 @@ def _update_policy_request( def _generate_budget_windows( budget_limits: Sequence[BudgetLimitEntry] | None, ) -> tuple[Mapping[str, object], ...] | None: - if budget_limits is None: + if not budget_limits: return None return tuple( MappingProxyType( @@ -402,15 +404,20 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite column_fields: Final = MappingProxyType( {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS} ) - request_metadata: Final = data.metadata or MappingProxyType({}) - folded_metadata: Final = {**request_metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict - columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates data_json in place + metadata: Final = data.metadata or MappingProxyType({}) + folded_metadata: Final = {**metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict + columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates in place expires: Final = ( now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None ) budget_reset_at: Final = ( get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None ) + key_rotation_at: Final = ( + now + timedelta(seconds=duration_in_seconds(duration=data.rotation_interval)) + if data.auto_rotate and data.rotation_interval + else None + ) return _verification_token_from_row( MappingProxyType( { @@ -418,6 +425,7 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite "metadata": encrypt_callback_vars(folded_metadata), "expires": expires, "budget_reset_at": budget_reset_at, + "key_rotation_at": key_rotation_at, "budget_limits": _generate_budget_windows(data.budget_limits), "object_permission": None, } @@ -3249,16 +3257,6 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) - await _enforce_custom_key_policy( - hook=_custom_key_policy_hook(proxy_server), - build_policy_request=lambda: _update_policy_request( - operation="update", - existing_key_row=existing_key_row, - non_default_values=non_default_values, - request=data, - ), - ) - # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) if new_key_alias != existing_key_row.key_alias: @@ -3278,6 +3276,16 @@ async def update_key_fn( existing_key_alias=existing_key_row.key_alias, ) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=data, + ), + ) + if prisma_client is None: raise Exception("Not connected to DB!") diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 35abb8fb718..63bbaa5ba4e 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict @@ -128,11 +128,18 @@ class BulkUpdateTeamKeysRequest(BaseModel): return self -CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"] +CustomKeyPolicyOperation: TypeAlias = Literal["generate", "update", "regenerate"] class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase): - """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request.""" + """What `general_settings.custom_key_policy` receives. + + `effective_key` is the verification token row as it will be written: the existing row overlaid with the + requested changes, with `duration` resolved to `expires` and `budget_duration` to `budget_reset_at`. Values the + proxy fills in after the policy stay at their defaults: `token`, `key_name`, `created_by`, `updated_by` and the + soft-budget `budget_id` on generate, the rotated token on regenerate, and the `object_permission` relation on + every operation (`object_permission_id` is set; read `request.object_permission` for the requested change). + """ model_config = ConfigDict(protected_namespaces=(), frozen=True) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 2b8c62523b9..b014e4428fd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( ResetSpendRequest, UpdateKeyRequest, ) +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -12362,7 +12363,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) received: list[CustomKeyPolicyRequest] = [] policy = _seven_day_policy(received) - data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + data = UpdateKeyRequest( + key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0, auto_rotate=True, rotation_interval="30d" + ) with ( patch( # test-quality-ok: cache eviction is outside the policy path @@ -12381,6 +12384,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp mock_prisma_client.update_data.assert_awaited_once() assert len(received) == 1 _assert_update_policy_request(received[0], data) + key_rotation_at = received[0].effective_key.key_rotation_at + assert key_rotation_at is not None + assert abs(key_rotation_at - (datetime.now(timezone.utc) + timedelta(days=30))) < timedelta(seconds=60) @pytest.mark.asyncio @@ -12695,6 +12701,23 @@ async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration assert effective_key.expires is None +def test_effective_key_after_update_swaps_the_object_permission_id_and_drops_the_stale_relation(): + existing_key = LiteLLM_VerificationToken( + token="tok", + object_permission_id="op-old", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-old", mcp_servers=["old"]), + ) + + effective_key = _effective_key_after_update( + existing_key_row=existing_key, non_default_values={"object_permission_id": "op-new"} + ) + + assert effective_key.object_permission_id == "op-new" + assert effective_key.object_permission is None + assert existing_key.object_permission is not None + assert existing_key.object_permission.mcp_servers == ["old"] + + def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it(): now = datetime(2026, 1, 1, tzinfo=timezone.utc) data = GenerateKeyRequest( @@ -12705,12 +12728,21 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat tags=["t1"], budget_duration="1d", max_budget=3.0, + budget_limits=[{"budget_duration": "1d", "max_budget": 5.0}], + auto_rotate=True, + rotation_interval="30d", + object_permission={"mcp_servers": ["srv"]}, key_type=LiteLLMKeyType.LLM_API, ) effective_key = _effective_key_for_generate(data=data, now=now) assert effective_key.expires == now + timedelta(days=5) + assert effective_key.key_rotation_at == now + timedelta(days=30) + assert effective_key.budget_limits is not None + assert effective_key.budget_limits[0]["max_budget"] == 5.0 + assert effective_key.budget_limits[0]["reset_at"] is not None + assert effective_key.object_permission is None assert effective_key.org_id == "org-1" assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]} assert effective_key.max_budget == 3.0 @@ -12722,6 +12754,18 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat assert data.guardrails == ["g1"] assert data.tags == ["t1"] assert data.duration == "5d" + assert data.budget_limits is not None + assert data.budget_limits[0].reset_at is None + assert data.object_permission is not None + assert data.object_permission.mcp_servers == ["srv"] + + +def test_effective_key_for_generate_stores_no_budget_windows_for_an_empty_list(): + effective_key = _effective_key_for_generate( + data=GenerateKeyRequest(budget_limits=[]), now=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + assert effective_key.budget_limits is None def test_effective_key_for_generate_without_duration_never_expires(): @@ -12731,6 +12775,7 @@ def test_effective_key_for_generate_without_duration_never_expires(): assert effective_key.expires is None assert effective_key.budget_reset_at is None + assert effective_key.key_rotation_at is None assert effective_key.key_type == "default" @@ -12764,6 +12809,27 @@ async def test_enforce_custom_key_policy_uses_the_default_denial_message(): assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule" +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_allows_when_the_decision_is_missing(): + received: list[CustomKeyPolicyRequest] = [] + + async def no_decision(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + return {} + + await _enforce_custom_key_policy(hook=no_decision, build_policy_request=_policy_request_for_generate) + + assert len(received) == 1 + assert received[0].operation == "generate" + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_never_builds_the_request_without_a_hook(): + await _enforce_custom_key_policy( + hook=None, build_policy_request=lambda: pytest.fail("policy request built without a hook") + ) + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ From 4b84c83788a8e9e4db02b0b85a5f43fe0c49ea30 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:59:23 -0700 Subject: [PATCH 051/297] fix(cost): split the cache read breakdown at the audio cache-read rate --- .../litellm_core_utils/llm_cost_calc/utils.py | 32 +++++++++++++++---- .../llm_cost_calc/test_llm_cost_calc_utils.py | 23 +++++++++++++ 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index dc689ca9618..8fc428b38ae 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1355,6 +1355,7 @@ class BilledTokenRates: input_cost_per_token: float output_cost_per_token: float cache_read_input_token_cost: float + cache_read_input_audio_token_cost: float cache_creation_input_token_cost: float cache_creation_input_token_cost_above_1hr: float output_cost_per_reasoning_token: float @@ -1366,6 +1367,7 @@ class BilledTokenRates: input_cost_per_token=self.input_cost_per_token * multiplier, output_cost_per_token=self.output_cost_per_token * multiplier, cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, + cache_read_input_audio_token_cost=self.cache_read_input_audio_token_cost * multiplier, cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, @@ -1389,15 +1391,16 @@ def _reasoning_token_count(usage: Usage) -> int: return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) -def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: - """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details - first, then the private top-level counters the Usage constructor mirrors cache tokens onto for - providers/callers that bypass the details.""" +def _cache_token_counts(usage: Usage) -> tuple[int, int, int, CacheCreationTokenDetails | None]: + """(cache read tokens, cached audio tokens, cache creation tokens, cache creation details): read from + prompt_tokens_details first, then the private top-level counters the Usage constructor mirrors cache + tokens onto for providers/callers that bypass the details.""" parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 return ( parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), + parsed["cache_hit_audio_tokens"] if parsed is not None else 0, parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), parsed["cache_creation_token_details"] if parsed is not None else None, ) @@ -1408,11 +1411,13 @@ def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRat cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" input_rate: Final = custom_cost_per_token["input_cost_per_token"] output_rate: Final = custom_cost_per_token["output_cost_per_token"] + cache_read_rate: Final = custom_cost_per_token.get("cache_read_input_token_cost", input_rate) cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate) return BilledTokenRates( input_cost_per_token=input_rate, output_cost_per_token=output_rate, - cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_read_input_token_cost=cache_read_rate, + cache_read_input_audio_token_cost=cache_read_rate, cache_creation_input_token_cost=cache_creation_rate, cache_creation_input_token_cost_above_1hr=cache_creation_rate, output_cost_per_reasoning_token=output_rate, @@ -1449,6 +1454,11 @@ def _cost_map_billed_rates( completion_base_cost=completion_base_cost, current_time=billing_time, ) + audio_cache_read_rate: Final = _get_cost_per_unit( + model_info, + _get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier), + None, + ) multiplier: Final = ( _get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(model_info, vertex_location) @@ -1458,6 +1468,9 @@ def _cost_map_billed_rates( input_cost_per_token=prompt_base_cost, output_cost_per_token=completion_base_cost, cache_read_input_token_cost=cache_read_cost_rate, + cache_read_input_audio_token_cost=( + audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost_rate + ), cache_creation_input_token_cost=cache_creation_cost_rate, cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, output_cost_per_reasoning_token=reasoning_rate, @@ -1530,7 +1543,9 @@ def get_token_type_cost_breakdown( if rates is None: return TokenTypeCostBreakdown(0.0, 0.0, 0.0) - cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage) + cache_read_tokens, cached_audio_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts( + usage + ) cache_creation_cost: Final = ( float(cache_creation_tokens) * rates.cache_creation_input_token_cost if custom_cost_per_token is not None @@ -1543,7 +1558,10 @@ def get_token_type_cost_breakdown( ) return TokenTypeCostBreakdown( reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, - cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost, + cache_read_cost=( + float(cache_read_tokens - cached_audio_tokens) * rates.cache_read_input_token_cost + + float(cached_audio_tokens) * rates.cache_read_input_audio_token_cost + ), cache_creation_cost=cache_creation_cost, rates=rates, ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 5ff1ab62698..21b96127a74 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4006,6 +4006,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp input_cost_per_token=6e-6, output_cost_per_token=3e-5, cache_read_input_token_cost=6e-7, + cache_read_input_audio_token_cost=6e-7, cache_creation_input_token_cost=7.5e-6, cache_creation_input_token_cost_above_1hr=0.0, output_cost_per_reasoning_token=3e-5, @@ -5252,3 +5253,25 @@ def test_cached_audio_tokens_billed_at_audio_cache_rate_through_model_info_looku prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai") assert prompt_cost == pytest.approx(300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7) + + +def test_cache_read_breakdown_splits_cached_audio_at_the_audio_cache_rate(_local_model_cost_map: None) -> None: + usage = Usage( + prompt_tokens=4863, + completion_tokens=1087, + total_tokens=5950, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=1693, + audio_tokens=3170, + cached_tokens=2816, + cached_tokens_details={"text_tokens": 896, "audio_tokens": 1920}, + ), + ) + + breakdown = get_token_type_cost_breakdown(model="gpt-realtime-2.1-mini", custom_llm_provider="openai", usage=usage) + prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai") + + assert breakdown.cache_read_cost == pytest.approx(896 * 6e-8 + 1920 * 3e-7) + assert breakdown.rates is not None + assert breakdown.rates.cache_read_input_audio_token_cost == pytest.approx(3e-7) + assert prompt_cost == pytest.approx((1693 - 896) * 6e-7 + (3170 - 1920) * 1e-5 + breakdown.cache_read_cost) From 8577d63ff5212173aa5b30bdfe8118b666c74179 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 00:14:49 +0000 Subject: [PATCH 052/297] fix(proxy): forward provider request id headers on mapped error responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 10 +++--- .../proxy/test_common_request_processing.py | 35 +++++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e6ed60ba177..9b8a98ca0e7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3445,15 +3445,13 @@ class ProxyBaseLLMRequestProcessing: # a failed request reports no timing, matching /v1/chat/completions read_timing_from_logging_obj=False, ) - # Extract headers from exception - check both e.headers and e.response.headers headers = getattr(e, "headers", None) or {} if not headers: - # Try to get headers from e.response.headers (httpx.Response) _response: Final = attribute_of(e, "response") - if _response is not None: - _response_headers: Final = getattr(_response, "headers", None) - if _response_headers: - headers = get_response_headers(dict(_response_headers)) + _response_headers: Final = getattr(_response, "headers", None) if _response is not None else None + _provider_headers: Final = _response_headers or getattr(e, "litellm_response_headers", None) + if _provider_headers: + headers = get_response_headers(dict(_provider_headers)) headers.update(custom_headers) # Call response headers hook for failure diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index efbb5eedad4..45e436f756a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8386,6 +8386,41 @@ async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500" +@pytest.mark.asyncio +async def test_handle_llm_api_exception_forwards_litellm_response_headers_when_response_is_synthetic(): + """Exception mapping hands the proxy a mapped error whose ``response`` is a synthetic empty + ``httpx.Response`` and parks the provider's real headers on ``litellm_response_headers``. + The client must still get the provider request id, as it does on a 200. + """ + import httpx + + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + mapped = litellm.BadRequestError( + message="OpenAIException - max_tokens is too large: 999999999.", + model="gpt-4o-mini", + llm_provider="openai", + ) + mapped.litellm_response_headers = httpx.Headers({"x-request-id": "req_openai_400"}) + assert dict(mapped.response.headers) == {} + + processor = ProxyBaseLLMRequestProcessing(data={}) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(ProxyException) as exc_info: + await processor._handle_llm_api_exception( + e=mapped, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.code == "400" + assert "max_tokens is too large: 999999999." in exc_info.value.message + assert exc_info.value.headers["llm_provider-x-request-id"] == "req_openai_400" + + class TestBackgroundResponseRetrievalGovernance: """LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines.""" From a159c7d98eb786a5bc7285dee0984512b6965a21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:38:23 -0700 Subject: [PATCH 053/297] fix(proxy): write the requested object permission row only after the key policy allows the update --- .../key_management_endpoints.py | 38 ++-- .../test_key_management_endpoints.py | 164 +++++++++++++++++- 2 files changed, 178 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 841c22da050..ad2ab58d357 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2464,12 +2464,6 @@ async def prepare_key_update_data( # sentinel for Json? columns, so store the JSON literal null non_default_values["budget_limits"] = json.dumps(None) - if "object_permission" in non_default_values: - non_default_values = await _handle_update_object_permission( - data_json=non_default_values, - existing_key_row=existing_key_row, - ) - _metadata: Final = existing_key_row.metadata or {} # validate model_max_budget @@ -2490,13 +2484,12 @@ async def prepare_key_update_data( async def _handle_update_object_permission( data_json: dict, existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, ) -> dict: - """ - Handle the update of object permission. - """ - from litellm.proxy.proxy_server import prisma_client + """Persist the requested object permission row and swap it for its id, only after the key policy allowed the write.""" + if "object_permission" not in data_json: + return data_json - # Use the common helper to handle the object permission update object_permission_id: Final = await handle_update_object_permission_common( data_json=data_json, existing_object_permission_id=existing_key_row.object_permission_id, @@ -2758,7 +2751,12 @@ async def _process_single_key_update( detail={"error": "Database not connected"}, ) - _data: Final = {**non_default_values, "token": update_key_request.key} + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) + _data: Final = {**update_values, "token": update_key_request.key} response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict "Mapping[str, object] | None", await prisma_client.update_data(token=update_key_request.key, data=_data), @@ -3289,18 +3287,23 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( await _update_key_row_with_soft_budget( prisma_client=prisma_client, key=key, data=data, - non_default_values=non_default_values, + non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, ) if "soft_budget" in data.model_fields_set - else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key})) + else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key})) ) # Delete - key from cache, since it's been updated! @@ -5324,7 +5327,12 @@ async def _execute_virtual_key_regeneration( request=data if data is not None else RegenerateKeyRequest(), ), ) - update_data.update(non_default_values) + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=key_in_db, + prisma_client=prisma_client, + ) + update_data.update(update_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index b014e4428fd..82458b5dfa3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from contextlib import ExitStack from typing import Final import json @@ -18,6 +19,7 @@ from litellm.proxy._types import ( GenerateKeyRequest, NewUserRequest, LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionBase, LiteLLM_OrganizationTable, LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj, @@ -1035,7 +1037,7 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_existing_permission(monkeypatch): +async def test_key_update_object_permissions_existing_permission(): """ Test updating object permissions when a key already has an existing object_permission_id. @@ -1055,9 +1057,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Mock existing key with object_permission_id existing_key_row = LiteLLM_VerificationToken( @@ -1097,6 +1097,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row, + prisma_client=mock_prisma_client, ) # Verify the object_permission was removed from data_json and object_permission_id was set @@ -1111,7 +1112,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_no_existing_permission(monkeypatch): +async def test_key_update_object_permissions_no_existing_permission(): """ Test creating object permissions when a key has no existing object_permission_id. @@ -1131,9 +1132,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_key_row_no_perm = LiteLLM_VerificationToken( token="test_token_hash_2", @@ -1164,6 +1163,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row_no_perm, + prisma_client=mock_prisma_client, ) # Verify new object_permission_id was set @@ -1174,7 +1174,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) @pytest.mark.asyncio -async def test_key_update_object_permissions_missing_permission_record(monkeypatch): +async def test_key_update_object_permissions_missing_permission_record(): """ Test creating object permissions when existing object_permission_id record is not found. @@ -1194,9 +1194,7 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_key_row_missing_perm = LiteLLM_VerificationToken( token="test_token_hash_3", @@ -1227,6 +1225,7 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row_missing_perm, + prisma_client=mock_prisma_client, ) # Verify new object_permission_id was set @@ -12470,6 +12469,153 @@ async def test_process_single_key_update_rejects_when_custom_key_policy_denies() assert [policy_request.operation for policy_request in received] == ["update"] +_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy" + + +def _record_object_permission_writes(mock_prisma_client: AsyncMock, events: list[str]) -> None: + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + + async def upsert(**_kwargs: object) -> MagicMock: + events.append("permission row upsert") + return MagicMock(object_permission_id=_OBJECT_PERMISSION_ID_AFTER_POLICY) + + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(side_effect=upsert) + + +def _recording_policy(events: list[str], allowed: bool): + async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + events.append("policy") + return {"decision": allowed, "message": "key max_budget must be 1000 or less"} + + return policy + + +def _assert_permission_row_written_after_policy(events: list[str], written: Mapping[str, object]) -> None: + assert events == ["policy", "permission row upsert"] + assert written["object_permission_id"] == _OBJECT_PERMISSION_ID_AFTER_POLICY + assert "object_permission" not in written + + +def _assert_permission_row_untouched(mock_prisma_client: AsyncMock, events: list[str]) -> None: + assert events == ["policy"] + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + + +def _update_with_object_permission(max_budget: float) -> UpdateKeyRequest: + return UpdateKeyRequest( + key=_POLICY_HASHED_TOKEN, + max_budget=max_budget, + object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"]), + ) + + +def _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed: bool) -> tuple[AsyncMock, list[str]]: + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", _recording_policy(events, allowed)) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", AsyncMock() + ) + return mock_prisma_client, events + + +async def _update_key_fn_with_object_permission(max_budget: float): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + return await update_key_fn( + request=MagicMock(), + data=_update_with_object_permission(max_budget=max_budget), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + +@pytest.mark.asyncio +async def test_update_key_fn_writes_the_object_permission_row_only_after_the_policy_allows(monkeypatch): + mock_prisma_client, events = _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed=True) + + await _update_key_fn_with_object_permission(max_budget=50.0) + + _assert_permission_row_written_after_policy(events, mock_prisma_client.update_data.await_args.kwargs["data"]) + + +@pytest.mark.asyncio +async def test_update_key_fn_denied_by_the_policy_leaves_the_object_permission_row_untouched(monkeypatch): + mock_prisma_client, events = _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed=False) + + with pytest.raises(ProxyException) as exc_info: + await _update_key_fn_with_object_permission(max_budget=5000.0) + + assert str(exc_info.value.code) == "403" + _assert_permission_row_untouched(mock_prisma_client, events) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_process_single_key_update_writes_the_object_permission_row_only_after_the_policy_allows(): + mock_prisma_client = AsyncMock() + updated_row = MagicMock() + updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + + await _process_single_key_update_under_policy( + mock_prisma_client, _update_with_object_permission(max_budget=50.0), _recording_policy(events, allowed=True) + ) + + _assert_permission_row_written_after_policy(events, mock_prisma_client.update_data.await_args.kwargs["data"]) + + +@pytest.mark.asyncio +async def test_process_single_key_update_denied_by_the_policy_leaves_the_object_permission_row_untouched(): + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy( + mock_prisma_client, _update_with_object_permission(max_budget=5000.0), _recording_policy(events, allowed=False) + ) + + assert exc_info.value.status_code == 403 + _assert_permission_row_untouched(mock_prisma_client, events) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_regenerate_writes_the_object_permission_row_only_after_the_policy_allows(): + mock_prisma_client = _make_regenerate_mock_prisma() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + data = RegenerateKeyRequest(max_budget=50.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + + with _regenerate_policy_mocks(_recording_policy(events, allowed=True), AsyncMock(), AsyncMock()): + await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + + _assert_permission_row_written_after_policy( + events, mock_prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"] + ) + + +@pytest.mark.asyncio +async def test_regenerate_denied_by_the_policy_leaves_the_object_permission_row_untouched(): + mock_prisma_client = _make_regenerate_mock_prisma() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + data = RegenerateKeyRequest(max_budget=5000.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + + with _regenerate_policy_mocks(_recording_policy(events, allowed=False), AsyncMock(), AsyncMock()): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + + assert exc_info.value.status_code == 403 + _assert_permission_row_untouched(mock_prisma_client, events) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + + @pytest.mark.asyncio async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys From f41c8556b5e22d1cd980df9ba2df7c0b8336d545 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 00:05:50 +0000 Subject: [PATCH 054/297] feat(jwt): allow virtual_key_claim_field per issuer Multi-IdP deployments can now set virtual_key_claim_field and unregistered_jwt_client_behavior on a JWTIssuerConfig entry. Tokens from that issuer use the issuer-specific claim path and no-match policy for the virtual key mapping lookup; issuers that omit them keep the global values. The auth flow now enters the mapping lookup when any issuer configures the field, not only when the global field is set. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 30 +++ litellm/proxy/auth/user_api_key_auth.py | 19 +- .../proxy/auth/test_user_api_key_auth.py | 225 +++++++++++++++++- tests/test_litellm/proxy/test__types.py | 58 +++++ 4 files changed, 324 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ae6c042ab3a..3ede847370d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4837,6 +4837,14 @@ class JWTIssuerConfig(BaseModel): default=None, description="Issuer-specific claim path to normalize into LiteLLM's end-user id.", ) + virtual_key_claim_field: str | None = Field( + default=None, + description="Issuer-specific claim path used for the virtual key mapping lookup. Falls back to the global field.", + ) + unregistered_jwt_client_behavior: UnregisteredJWTClientBehavior | None = Field( + default=None, + description="Issuer-specific policy when the virtual key claim has no mapping. Falls back to the global policy.", + ) model_config = { "extra": "forbid", @@ -5063,6 +5071,28 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): super().__init__(**kwargs) + def get_issuer_config(self, issuer: str | None) -> JWTIssuerConfig | None: + if issuer is None or self.issuers is None: + return None + return next((config for config in self.issuers if config.issuer == issuer), None) + + def is_virtual_key_mapping_configured(self) -> bool: + if self.virtual_key_claim_field is not None: + return True + return any(config.virtual_key_claim_field is not None for config in self.issuers or ()) + + def get_virtual_key_claim_field(self, issuer: str | None) -> str | None: + issuer_config: Final = self.get_issuer_config(issuer) + if issuer_config is not None and issuer_config.virtual_key_claim_field is not None: + return issuer_config.virtual_key_claim_field + return self.virtual_key_claim_field + + def get_unregistered_jwt_client_behavior(self, issuer: str | None) -> UnregisteredJWTClientBehavior: + issuer_config: Final = self.get_issuer_config(issuer) + if issuer_config is not None and issuer_config.unregistered_jwt_client_behavior is not None: + return issuer_config.unregistered_jwt_client_behavior + return self.unregistered_jwt_client_behavior + class PrismaCompatibleUpdateDBModel(TypedDict, total=False): model_name: str diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9828311112e..f1b373da898 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -987,9 +987,12 @@ async def _resolve_jwt_to_virtual_key( - Raises HTTPException: REJECT policy hit, missing claim under REJECT/AUTO_REGISTER, or other policy violations. """ - virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.virtual_key_claim_field + raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM) + normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None + virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.get_virtual_key_claim_field(normalized_issuer) if virtual_key_claim_field is None: return None + behavior: Final = jwt_handler.litellm_jwtauth.get_unregistered_jwt_client_behavior(normalized_issuer) claim_value: Final = get_nested_value( data=jwt_claims, @@ -1006,7 +1009,6 @@ async def _resolve_jwt_to_virtual_key( # simply by presenting a JWT that omits the configured field. For # AUTO_REGISTER there is no stable identity to map without a claim # value, so we deny rather than create a sentinel-keyed record. - behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior if behavior in ( UnregisteredJWTClientBehavior.REJECT, UnregisteredJWTClientBehavior.AUTO_REGISTER, @@ -1021,7 +1023,13 @@ async def _resolve_jwt_to_virtual_key( return None cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value)) - cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key) + raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key) + sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER + cached_mapping: Final = ( + None + if raw_cached_mapping == _JWT_PROXY_ADMIN_SENTINEL and not sentinel_written_by_this_policy + else raw_cached_mapping + ) if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL: # Previously resolved to a proxy admin via auth_builder; skip the @@ -1030,7 +1038,6 @@ async def _resolve_jwt_to_virtual_key( return None if cached_mapping == "__NO_MAPPING__": - behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior if behavior == UnregisteredJWTClientBehavior.REJECT: raise HTTPException( status_code=403, @@ -1093,8 +1100,6 @@ async def _resolve_jwt_to_virtual_key( ) # No mapping found (DB miss or no DB) — apply no-match policy. - behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior - if behavior == UnregisteredJWTClientBehavior.REJECT: # Cache the miss before raising so repeated rejections are served from # cache and don't re-query the DB on every request. @@ -1428,7 +1433,7 @@ async def _user_api_key_auth_builder( # unnecessary DB queries in auth_builder do_standard_jwt_auth = True pending_auto_register: _PendingAutoRegister | None = None - if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None: + if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured(): # Decode JWT to get claims without running full auth_builder jwt_claims: dict | None if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not is_jwt: diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 0cdbcde6abc..dac6b7aedc5 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -13,7 +13,7 @@ from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest -from fastapi import status +from fastapi import HTTPException, status import litellm import litellm.proxy.proxy_server @@ -7442,3 +7442,226 @@ async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer) assert data["model"] == ("foo" if layer == "unclaimed" else encoded) await _normalize_claude_model(data, token, request, "/v1/messages") assert data["model"] == ("foo" if layer == "unclaimed" else encoded) + + +ISSUER_ONE = "https://issuer-one.example.com" +ISSUER_TWO = "https://issuer-two.example.com" + + +def _per_issuer_virtual_key_jwt_handler( + global_claim_field: str | None, global_behavior: str = "fallback_team_mapping" +) -> MagicMock: + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field=global_claim_field, + unregistered_jwt_client_behavior=global_behavior, + issuers=[ + { + "issuer": ISSUER_ONE, + "jwks_url": f"{ISSUER_ONE}/keys", + "audience": "audience-one", + "team_id_jwt_field": "sub", + }, + { + "issuer": ISSUER_TWO, + "jwks_url": f"{ISSUER_TWO}/keys", + "audience": "audience-two", + "virtual_key_claim_field": "sub", + "unregistered_jwt_client_behavior": "reject", + }, + ], + ) + return jwt_handler + + +def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]: + find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token)) + prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first))) + return prisma_client, find_first + + +def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]: + return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True} + + +@pytest.mark.asyncio +async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for_the_db_lookup(): + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None) + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping("hashed-mapped-key") + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="hashed-mapped-key", + value=UserAPIKeyAuth(token="hashed-mapped-key", api_key="hashed-mapped-key", team_id="svc-team"), + ) + + resolved = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "svc-account-7"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7")) + assert isinstance(resolved, UserAPIKeyAuth) + assert resolved.token == "hashed-mapped-key" + assert resolved.team_id == "svc-team" + assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key" + + +@pytest.mark.asyncio +async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer(): + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None) + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) + + team_issuer_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert team_issuer_result is None + find_first.assert_not_awaited() + + with pytest.raises(HTTPException) as exc: + await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "unknown-svc"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail) + find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc")) + + +@pytest.mark.asyncio +async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_reject(): + from litellm.proxy.auth.user_api_key_auth import _JWT_PROXY_ADMIN_SENTINEL, _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register") + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL) + + auto_register_issuer_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert auto_register_issuer_result is None + find_first.assert_not_awaited() + + with pytest.raises(HTTPException) as exc: + await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "admin-7"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "No registered mapping for sub='admin-7'" in str(exc.value.detail) + find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7")) + + +@pytest.mark.asyncio +async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_field(): + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="client_id") + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) + + with_claim = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha", "client_id": "app-9"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + without_claim = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert with_claim is None + assert without_claim is None + find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9")) + + +@pytest.mark.asyncio +async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configures_the_claim_field(): + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtYWNjb3VudC03In0.signature" + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None) + jwt_handler.auth_jwt = AsyncMock( + return_value={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "svc-account-7"} + ) + mapped_key = UserAPIKeyAuth(token="hashed-mapped-key", api_key="hashed-mapped-key", team_id="svc-team") + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam + "litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True} + ), + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.user_api_key_cache", DualCache() + ), + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.proxy_logging_obj", MagicMock() + ), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: the regression is whether the builder reaches this seam at all + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=mapped_key, + ) as resolve_mock, + patch( # test-quality-ok: a mapped key must short-circuit standard JWT auth; reaching it is the failure + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + side_effect=AssertionError("standard JWT auth must not run for a mapped virtual key"), + ), + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o-mini"}, + ) + + resolve_mock.assert_awaited_once() + assert resolve_mock.await_args.kwargs["jwt_claims"][JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == ISSUER_TWO + assert result.api_key == "hashed-mapped-key" + assert result.team_id == "svc-team" diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 26bb1533da4..9e1486ce90f 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -277,3 +277,61 @@ def test_team_membership_budget_table_present_still_works(): } result = LiteLLM_TeamMembership.model_validate(data) assert result.litellm_budget_table is None + + +def test_a_jwt_issuer_can_override_the_virtual_key_claim_field_while_other_issuers_keep_the_global_one(): + from litellm.proxy._types import LiteLLM_JWTAuth, UnregisteredJWTClientBehavior + + jwt_auth = LiteLLM_JWTAuth( + virtual_key_claim_field="client_id", + issuers=[ + { + "issuer": "https://team-idp.example.com", + "jwks_url": "https://team-idp.example.com/keys", + "audience": "litellm", + "team_id_jwt_field": "sub", + }, + { + "issuer": "https://service-idp.example.com", + "jwks_url": "https://service-idp.example.com/keys", + "audience": "litellm", + "virtual_key_claim_field": "sub", + "unregistered_jwt_client_behavior": "reject", + }, + ], + ) + + assert jwt_auth.get_virtual_key_claim_field("https://service-idp.example.com") == "sub" + assert jwt_auth.get_unregistered_jwt_client_behavior("https://service-idp.example.com") is ( + UnregisteredJWTClientBehavior.REJECT + ) + assert jwt_auth.get_virtual_key_claim_field("https://team-idp.example.com") == "client_id" + assert jwt_auth.get_unregistered_jwt_client_behavior("https://team-idp.example.com") is ( + UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING + ) + assert jwt_auth.get_virtual_key_claim_field(None) == "client_id" + assert jwt_auth.get_virtual_key_claim_field("https://unknown-idp.example.com") == "client_id" + + +@pytest.mark.parametrize( + ("global_field", "issuer_field", "is_configured"), + ((None, None, False), ("sub", None, True), (None, "sub", True)), +) +def test_virtual_key_mapping_counts_as_configured_when_any_issuer_sets_the_claim_field( + global_field, issuer_field, is_configured +): + from litellm.proxy._types import LiteLLM_JWTAuth + + jwt_auth = LiteLLM_JWTAuth( + virtual_key_claim_field=global_field, + issuers=[ + { + "issuer": "https://idp.example.com", + "jwks_url": "https://idp.example.com/keys", + "audience": "litellm", + "virtual_key_claim_field": issuer_field, + } + ], + ) + + assert jwt_auth.is_virtual_key_mapping_configured() is is_configured From a94c060b84815438dc3a6dda379e0fb7fa60448f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:48:00 -0700 Subject: [PATCH 055/297] fix(cost): fill the missing realtime cache-read rates --- ...odel_prices_and_context_window_backup.json | 9 ++++-- model_prices_and_context_window.json | 9 ++++-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 28 +++++++++++++++++++ 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index de9867f9eee..2ce8914ba5a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5513,7 +5513,8 @@ }, "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5546,7 +5547,8 @@ }, "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5683,6 +5685,7 @@ }, "azure/gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -5715,6 +5718,7 @@ }, "azure/gpt-realtime-mini-2025-10-06": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -32683,6 +32687,7 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index de9867f9eee..2ce8914ba5a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5513,7 +5513,8 @@ }, "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5546,7 +5547,8 @@ }, "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5683,6 +5685,7 @@ }, "azure/gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -5715,6 +5718,7 @@ }, "azure/gpt-realtime-mini-2025-10-06": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -32683,6 +32687,7 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 21b96127a74..7cd19c65887 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5275,3 +5275,31 @@ def test_cache_read_breakdown_splits_cached_audio_at_the_audio_cache_rate(_local assert breakdown.rates is not None assert breakdown.rates.cache_read_input_audio_token_cost == pytest.approx(3e-7) assert prompt_cost == pytest.approx((1693 - 896) * 6e-7 + (3170 - 1920) * 1e-5 + breakdown.cache_read_cost) + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_prompt_cost"), + ( + pytest.param("azure/gpt-realtime-2025-08-28", "azure", 300 * 4e-6 + 100 * 4e-7 + 200 * 3.2e-5 + 400 * 4e-7, id="azure-gpt-realtime"), + pytest.param("azure/gpt-realtime-1.5-2026-02-23", "azure", 300 * 4e-6 + 100 * 4e-7 + 200 * 3.2e-5 + 400 * 4e-7, id="azure-gpt-realtime-1.5"), + pytest.param("azure/gpt-realtime-mini", "azure", 300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7, id="azure-gpt-realtime-mini"), + pytest.param("gpt-realtime-mini", "openai", 300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7, id="openai-gpt-realtime-mini"), + ), +) +def test_realtime_models_bill_cached_text_and_audio_at_their_cache_read_rates( + _local_model_cost_map: None, model: str, custom_llm_provider: str, expected_prompt_cost: float +) -> None: + usage = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=400, + audio_tokens=600, + cached_tokens=500, + cached_tokens_details={"text_tokens": 100, "audio_tokens": 400}, + ), + ) + + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=custom_llm_provider) + assert prompt_cost == pytest.approx(expected_prompt_cost) From 555e321cf170ad2d15a2932fccdea104da1c39b6 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 01:05:51 +0000 Subject: [PATCH 056/297] fix(router): record flat retry attempts and cap retries from attempted_retries Router.log_retry used to copy the failed attempt's kwargs and metadata into metadata.previous_models. Nothing downstream read those copies, but they carried client credentials into spend logs and grew the payload on every retry. Each attempt now leaves a flat record (model group, deployment id, exception type and string, attempt number), which drops RETRY_BREADCRUMB_EXCLUDED_KWARGS and the per-retry credential masking. num_retries_per_request was enforced from len(previous_models), which only looked at the metadata bucket and never exceeded four records. The sync and async client wrappers and the Rust lifecycle guard now read attempted_retries from whichever metadata bucket the call carries. Resolves LIT-7505 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 13 +++ litellm/router.py | 45 +++----- litellm/rust_bridge/lifecycle.py | 16 +-- litellm/types/router.py | 8 ++ litellm/utils.py | 18 +-- .../test_router_helper_utils.py | 29 +++-- .../rust_bridge/test_lifecycle.py | 30 +++++ tests/test_litellm/test_router.py | 106 +++++++++++------- tests/test_litellm/test_utils.py | 47 ++++++++ tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- 11 files changed, 209 insertions(+), 107 deletions(-) create mode 100644 tests/test_litellm/rust_bridge/test_lifecycle.py diff --git a/litellm/__init__.py b/litellm/__init__.py index ccfbf80369f..261457d6889 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries) +num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop ####### SECRET MANAGERS ##################### secret_manager_client: Optional[Any] = ( None # list of instantiated key management clients - e.g. azure kv, infisical, etc. diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index aa7d6ca1699..ecd9cdac88b 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -303,6 +303,19 @@ def get_metadata_variable_name_from_kwargs( return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" +def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool: + """ + Whether the Router retry about to run (``attempted_retries`` >= 1 in the metadata bucket) is past the cap + """ + if num_retries_per_request is None: + return False + metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) + if not isinstance(metadata, Mapping): + return False + attempted_retries: Final = metadata.get("attempted_retries") + return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries + + def get_or_create_metadata_bucket( request_data: dict, ) -> tuple[Literal["metadata", "litellm_metadata"], dict]: diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..a46ffa83b05 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -96,7 +96,6 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, - mask_credentials_in_payload, mask_sensitive_structure, ) from litellm.litellm_core_utils.token_counter import offload_token_count @@ -242,6 +241,7 @@ from litellm.types.router import ( ModelGroupInfo, OptionalPreCallChecks, PreRoutingStrategy, + RetryAttemptRecord, RetryPolicy, RouterCacheEnum, RouterErrors, @@ -623,20 +623,6 @@ def _replay_live_router_model_cost() -> None: set_live_deployment_replay(_replay_live_router_model_cost) -# Kwargs that carry no signal about the failed attempt, so log_retry drops them from a -# breadcrumb entirely: the request payload, the proxy's snapshot of the inbound request (its body -# aliases the live request metadata, earlier breadcrumbs included, so copying it would nest every -# breadcrumb inside the next one), and the router-internal walk state. Credentials are handled -# separately by mask_credentials_in_payload, which scrubs credential-named values from whatever -# kwargs remain rather than trying to enumerate every credential-bearing key here. -RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( - ( - "messages", - "original_function", - "attempted_targets", - "proxy_server_request", - ) -) RETRY_BREADCRUMB_LIMIT: Final = 4 @@ -8374,31 +8360,28 @@ class Router: def log_retry(self, kwargs: dict, e: Exception) -> dict: """ - When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing + When a retry or fallback happens, record which model group, deployment and attempt just failed and why """ _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var] - attempt_kwargs: Final = MappingProxyType( - {k: v for k, v in kwargs.items() if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS} - ) - attempt_metadata: Final = MappingProxyType( - {k: v for k, v in request_metadata.items() if k != "previous_models"} - ) - previous_model: Final = MappingProxyType( - { - "exception_type": type(e).__name__, - "exception_string": str(e), - **attempt_kwargs, - _metadata_var: attempt_metadata, - } - ) + model_group: Final = kwargs.get("model") + model_info: Final = request_metadata.get("model_info") + deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None + attempted_retries: Final = request_metadata.get("attempted_retries") + attempt_record: Final[RetryAttemptRecord] = { + "model_group": model_group if isinstance(model_group, str) else None, + "deployment_id": deployment_id if isinstance(deployment_id, str) else None, + "exception_type": type(e).__name__, + "exception_string": str(e), + "attempted_retries": attempted_retries if type(attempted_retries) is int else None, + } earlier_breadcrumbs: Final = request_metadata.get("previous_models") kept_breadcrumbs: Final[tuple[object, ...]] = ( tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] if isinstance(earlier_breadcrumbs, (list, tuple)) else () ) - breadcrumbs: Final = (*kept_breadcrumbs, mask_credentials_in_payload(previous_model)) + breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict return kwargs diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index f5e0c1b0fc6..f1cc912129d 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -99,23 +99,13 @@ def setup( def check_limits(kwargs: Mapping[str, object]) -> None: import litellm + from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor if litellm.max_budget and current_cost > litellm.max_budget: raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) - metadata: Final = kwargs.get("metadata") - if isinstance(metadata, Mapping): - typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata - Mapping[str, object], metadata - ) - previous: Final = typed_metadata.get("previous_models") - if ( - isinstance(previous, list) - and litellm.num_retries_per_request is not None - and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history - >= litellm.num_retries_per_request - ): - raise RuntimeError("Max retries per request hit!") + if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + raise RuntimeError("Max retries per request hit!") def finalize( diff --git a/litellm/types/router.py b/litellm/types/router.py index fc09c40fe08..fecc0e00f99 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -883,6 +883,14 @@ class RouterModelGroupAliasItem(TypedDict): hidden: bool # if 'True', don't return on `.get_model_list` +class RetryAttemptRecord(TypedDict): + model_group: ReadOnly[str | None] + deployment_id: ReadOnly[str | None] + exception_type: ReadOnly[str] + exception_string: ReadOnly[str] + attempted_retries: ReadOnly[int | None] + + VALID_LITELLM_ENVIRONMENTS = [ "development", "staging", diff --git a/litellm/utils.py b/litellm/utils.py index 394ab4b4094..98881e68986 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -81,7 +81,7 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit, normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, ) @@ -1509,12 +1509,8 @@ def client(original_function): call_type = original_function.__name__ if _is_async_request(kwargs): # [OPTIONAL] CHECK MAX RETRIES / REQUEST - if litellm.num_retries_per_request is not None: - # check if previous_models passed in as ['litellm_params']['metadata]['previous_models'] - previous_models = (kwargs.get("metadata") or {}).get("previous_models", None) - if previous_models is not None: - if litellm.num_retries_per_request <= len(previous_models): - raise Exception("Max retries per request hit!") + if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + raise Exception("Max retries per request hit!") # MODEL CALL result = original_function(*args, **kwargs) @@ -1573,12 +1569,8 @@ def client(original_function): ) # [OPTIONAL] CHECK MAX RETRIES / REQUEST - if litellm.num_retries_per_request is not None: - # check if previous_models passed in as ['litellm_params']['metadata]['previous_models'] - previous_models = (kwargs.get("metadata") or {}).get("previous_models", None) - if previous_models is not None: - if litellm.num_retries_per_request <= len(previous_models): - raise Exception("Max retries per request hit!") + if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + raise Exception("Max retries per request hit!") # [OPTIONAL] CHECK CACHE print_verbose( diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 7dbac243d55..5b06c5fdb01 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1,3 +1,4 @@ +import json import os import traceback from dotenv import load_dotenv @@ -628,17 +629,29 @@ def test_deployment_callback_respects_cooldown_time(model_list): assert mock_set.call_args.kwargs["time_to_cooldown"] == 0 -def test_log_retry(model_list): - """Test if the '_log_retry' function is working correctly""" - import time - +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_log_retry(model_list, metadata_key): + """log_retry appends one flat record per failed attempt and copies neither the request kwargs nor + the request metadata into it""" router = Router(model_list=model_list) new_kwargs = router.log_retry( - kwargs={"metadata": {}}, - e=Exception(), + kwargs={ + "model": "gpt-3.5-turbo", + "api_key": "sk-must-not-be-recorded", + "messages": [{"role": "user", "content": "hi"}], + metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"}, + }, + e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"), ) - assert "metadata" in new_kwargs - assert "previous_models" in new_kwargs["metadata"] + assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [ + { + "model_group": "gpt-3.5-turbo", + "deployment_id": "deployment-1", + "exception_type": "RateLimitError", + "exception_string": "litellm.RateLimitError: slow down", + "attempted_retries": 2, + } + ] def test_update_usage(model_list): diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/test_litellm/rust_bridge/test_lifecycle.py new file mode 100644 index 00000000000..1f0b5591c2b --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_lifecycle.py @@ -0,0 +1,30 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.lifecycle import check_limits + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize( + "cap, attempted_retries, refused", + [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)], + ids=[ + "cap-above-four-reached", + "cap-above-four-not-reached", + "first-attempt-passes-cap-of-zero", + "cap-of-zero-refuses-first-retry", + ], +) +def test_check_limits_reads_attempted_retries( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, attempted_retries: int, refused: bool +) -> None: + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + monkeypatch.setattr(litellm, "max_budget", None) + kwargs: Final = {"model": "mistral/mistral-ocr-latest", metadata_key: {"attempted_retries": attempted_retries}} + if refused: + with pytest.raises(RuntimeError, match="Max retries per request hit!"): + check_limits(kwargs) + else: + check_limits(kwargs) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f5e9b2091a0..6150ce287ed 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10607,6 +10607,7 @@ def _cyclic_fallback_router(num_retries=0): "api_key": "sk-fake", "mock_response": "litellm.InternalServerError", }, + "model_info": {"id": f"{group}-deployment"}, } for group in groups ], @@ -10656,28 +10657,37 @@ async def test_cyclic_fallback_graph_does_not_amplify_one_request(): assert sum(len(message) for message in capture.messages) < 5_000 +_FLAT_ATTEMPT_RECORD_KEYS = frozenset( + {"model_group", "deployment_id", "exception_type", "exception_string", "attempted_retries"} +) +_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip" + + @pytest.mark.asyncio -async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): - """log_retry copies every kwarg into previous_models, which reaches spend logs and - logging callbacks. The set of already-attempted groups is router-internal walk state - with no diagnostic value there, and it is the one entry that is not a plain scalar. - A retry has to be configured for the walk state to reach log_retry at all.""" +async def test_retry_records_are_flat_and_name_the_failed_group_on_fallback_hops(): + """Each failed attempt leaves a flat record in previous_models, which reaches spend logs and + logging callbacks. Nothing downstream reads the failed attempt's kwargs or metadata, and copying + them is what carried client credentials and multiplied the payload on every retry. A fallback hop + calls log_retry too, so the record has to name the group that failed, not the one taken next.""" router = _cyclic_fallback_router(num_retries=1) capture = _LogCapture(logging.ERROR) recorder = _FallbackAttemptRecorder() await _drive_cyclic_fallback(router, capture, recorder) - breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] - assert breadcrumbs, "no retry breadcrumbs were recorded" - assert any( - "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs - ), "no breadcrumb carried router walk state, so this test cannot see the leak" - for breadcrumb in breadcrumbs: - assert "attempted_targets" not in breadcrumb - - -_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip" + records = [record for hop in recorder.breadcrumbs_per_target for record in hop] + assert records, "no retry records were recorded" + for record in records: + assert set(record) == _FLAT_ATTEMPT_RECORD_KEYS + assert record["exception_type"] == "InternalServerError" + assert record["deployment_id"] == f"{record['model_group']}-deployment" + group_failed_before_hop = {"group-b": "group-a", "group-c": "group-b", "group-d": "group-c"} + for failed_target, hop_records in zip(recorder.failed_targets, recorder.breadcrumbs_per_target): + groups = [record["model_group"] for record in hop_records] + first_own_attempt = groups.index(failed_target) + assert groups[first_own_attempt - 1] == group_failed_before_hop[failed_target] + assert set(groups[first_own_attempt:]) == {failed_target} + assert [record["attempted_retries"] for record in hop_records[first_own_attempt:]][:2] == [0, 1] @pytest.mark.parametrize( @@ -10703,22 +10713,20 @@ _BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doN ], ) @pytest.mark.asyncio -async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_key, request_kwargs): - """log_retry copies kwargs into previous_models, which reaches spend logs and logging callbacks. - Any of these kwargs can carry a client's forwarded Authorization token or a provider key, and a - breadcrumb has no diagnostic use for the raw secret. A denylist of key names is always one new - credential kwarg behind, so log_retry scrubs credential-named values by pattern instead: the - container still reaches the breadcrumb, but the raw secret never does, whatever key holds it.""" +async def test_retry_records_never_carry_a_forwarded_credential(container_key, request_kwargs): + """previous_models reaches spend logs and logging callbacks. Any request kwarg can carry a client's + forwarded Authorization token or a provider key, so the record must not carry request kwargs at + all: neither the credential-bearing container nor the raw secret, whatever key holds it.""" router = _cyclic_fallback_router(num_retries=1) capture = _LogCapture(logging.ERROR) metadata = {} await _drive_cyclic_fallback(router, capture, metadata=metadata, **request_kwargs) - breadcrumbs = metadata["previous_models"] - assert breadcrumbs, "no retry breadcrumbs were recorded" - dumped = json.dumps(breadcrumbs, default=str) - assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + records = metadata["previous_models"] + assert records, "no retry records were recorded" + dumped = json.dumps(records) + assert container_key not in dumped assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -10743,7 +10751,7 @@ async def _fail_one_proxy_shaped_request(router, request_marker): shallow copy of the request, so body["metadata"] is the very same dict the router later stamps previous_models onto.""" metadata = {"request_marker": request_marker} - with pytest.raises(litellm.InternalServerError): + with pytest.raises((litellm.InternalServerError, litellm.APIConnectionError)): await router.acompletion( model="broken-group", messages=[{"role": "user", "content": "hi"}], @@ -10769,34 +10777,52 @@ def _nested_breadcrumb_lists(node): @pytest.mark.asyncio -async def test_retry_breadcrumbs_stay_per_request_and_flat_across_failing_requests(): - """Every failed attempt appends a breadcrumb to metadata["previous_models"], and the proxy's +async def test_retry_records_stay_per_request_and_flat_across_failing_requests(): + """Every failed attempt appends a record to metadata["previous_models"], and the proxy's request snapshot aliases that same metadata dict. Kept on the Router and copied wholesale, - each breadcrumb embedded every earlier one from every earlier request, so the breadcrumb + each breadcrumb once embedded every earlier one from every earlier request, so the breadcrumb tree, and with it the debug repr of the kwargs, roughly doubled on each failed attempt until a single-worker proxy spent minutes in the redaction regex and stopped answering.""" router = _always_failing_router(num_retries=2) - breadcrumbs_per_request = [ + records_per_request = [ await _fail_one_proxy_shaped_request(router, f"request-{request_number}") for request_number in range(1, 7) ] - for request_number, breadcrumbs in enumerate(breadcrumbs_per_request, start=1): - assert len(breadcrumbs) == 3, "one initial attempt plus two retries failed, each leaving one breadcrumb" - assert {breadcrumb["metadata"]["request_marker"] for breadcrumb in breadcrumbs} == {f"request-{request_number}"} - for breadcrumb in breadcrumbs: - assert _nested_breadcrumb_lists(breadcrumb) == [] - assert len({len(repr(breadcrumbs)) for breadcrumbs in breadcrumbs_per_request}) == 1 + for records in records_per_request: + assert [record["attempted_retries"] for record in records] == [0, 1, 2] + for record in records: + assert set(record) == _FLAT_ATTEMPT_RECORD_KEYS + assert _nested_breadcrumb_lists(record) == [] + assert len({len(repr(records)) for records in records_per_request}) == 1 @pytest.mark.asyncio -async def test_retry_breadcrumbs_keep_only_the_last_four_attempts(): +async def test_retry_records_keep_only_the_last_four_attempts(): router = _always_failing_router(num_retries=6) - breadcrumbs = await _fail_one_proxy_shaped_request(router, "request-1") + records = await _fail_one_proxy_shaped_request(router, "request-1") - assert len(breadcrumbs) == 4 - assert [breadcrumb["metadata"]["attempted_retries"] for breadcrumb in breadcrumbs] == [3, 4, 5, 6] + assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6] + + +@pytest.mark.asyncio +async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypatch): + """The cap used to be read off len(previous_models), which never exceeds four, so any cap above + four was inert. Reading the Router's attempted_retries counter instead lets a cap of five refuse + retries five and six before they reach the deployment.""" + monkeypatch.setattr(litellm, "num_retries_per_request", 5) + router = _always_failing_router(num_retries=6) + + records = await _fail_one_proxy_shaped_request(router, "request-1") + + assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6] + assert ["Max retries per request hit!" in record["exception_string"] for record in records] == [ + False, + False, + True, + True, + ] @pytest.mark.asyncio diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 835e87aff88..523feb2e54a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4061,6 +4061,53 @@ class TestMetadataNoneHandling: assert metadata == {} +_RETRY_CAP_CASES: Final = ( + pytest.param(5, {"attempted_retries": 5}, True, id="cap-above-four-reached"), + pytest.param(5, {"attempted_retries": 4}, False, id="cap-above-four-not-reached"), + pytest.param(0, {"attempted_retries": 0}, False, id="first-attempt-passes-cap-of-zero"), + pytest.param(0, {"attempted_retries": 1}, True, id="cap-of-zero-refuses-first-retry"), + pytest.param(5, {"previous_models": ("a", "b", "c", "d", "e")}, False, id="breadcrumb-count-is-not-the-cap"), + pytest.param(5, None, False, id="metadata-none"), +) + + +def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, object]: + return { + "model": "openai/gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "api_key": "sk-fake", + "mock_response": "ok", + metadata_key: metadata, + } + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) +def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused): + """num_retries_per_request is enforced from the Router's attempted_retries counter in whichever + metadata bucket the call carries, so callers on litellm_metadata and caps above four both work""" + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) + if refused: + with pytest.raises(Exception, match="Max retries per request hit!"): + litellm.completion(**kwargs) + else: + assert litellm.completion(**kwargs).choices[0].message.content == "ok" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) +async def test_num_retries_per_request_reads_attempted_retries_async(monkeypatch, metadata_key, cap, metadata, refused): + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) + if refused: + with pytest.raises(Exception, match="Max retries per request hit!"): + await litellm.acompletion(**kwargs) + else: + assert (await litellm.acompletion(**kwargs)).choices[0].message.content == "ok" + + class TestValidateAndFixThinkingParam: """Tests for validate_and_fix_thinking_param.""" diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index e7ebc5b3018..77d9ef167d0 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -806,7 +806,7 @@ async def test_shared_call_limits_still_reject_before_reading_ocr_file( monkeypatch.setattr(litellm, "_current_cost", 2) monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None) expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError - arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"previous_models": ["earlier"]}} + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"attempted_retries": 1}} with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"): await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) assert reads == [] From 566da87771b65665d9d8489898a86db693b5ca06 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 01:20:46 +0000 Subject: [PATCH 057/297] test(router): expect the exact error per retry-cap case and drop explanatory docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 3 --- tests/test_litellm/test_router.py | 9 +++------ tests/test_litellm/test_utils.py | 2 -- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index ecd9cdac88b..6e76bf9d49e 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -304,9 +304,6 @@ def get_metadata_variable_name_from_kwargs( def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool: - """ - Whether the Router retry about to run (``attempted_retries`` >= 1 in the metadata bucket) is past the cap - """ if num_retries_per_request is None: return False metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6150ce287ed..eb29f717a20 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10746,12 +10746,12 @@ def _always_failing_router(num_retries): ) -async def _fail_one_proxy_shaped_request(router, request_marker): +async def _fail_one_proxy_shaped_request(router, request_marker, expected_error=litellm.InternalServerError): """The proxy hands the router a metadata dict and a proxy_server_request whose body is a shallow copy of the request, so body["metadata"] is the very same dict the router later stamps previous_models onto.""" metadata = {"request_marker": request_marker} - with pytest.raises((litellm.InternalServerError, litellm.APIConnectionError)): + with pytest.raises(expected_error): await router.acompletion( model="broken-group", messages=[{"role": "user", "content": "hi"}], @@ -10808,13 +10808,10 @@ async def test_retry_records_keep_only_the_last_four_attempts(): @pytest.mark.asyncio async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypatch): - """The cap used to be read off len(previous_models), which never exceeds four, so any cap above - four was inert. Reading the Router's attempted_retries counter instead lets a cap of five refuse - retries five and six before they reach the deployment.""" monkeypatch.setattr(litellm, "num_retries_per_request", 5) router = _always_failing_router(num_retries=6) - records = await _fail_one_proxy_shaped_request(router, "request-1") + records = await _fail_one_proxy_shaped_request(router, "request-1", expected_error=litellm.APIConnectionError) assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6] assert ["Max retries per request hit!" in record["exception_string"] for record in records] == [ diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 523feb2e54a..3f3b8ce5343 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4084,8 +4084,6 @@ def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused): - """num_retries_per_request is enforced from the Router's attempted_retries counter in whichever - metadata bucket the call carries, so callers on litellm_metadata and caps above four both work""" monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: From 0429c645043e363ae2a87700a858349657850f6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:22:31 -0700 Subject: [PATCH 058/297] fix(proxy): keep an empty duration out of the legacy update hook on key regenerate --- .../key_management_endpoints.py | 7 ++- .../test_key_management_endpoints.py | 57 ++++++++++++++++++- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ad2ab58d357..8aaaba09e22 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -433,12 +433,17 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite ) +_EMPTY_DURATION_MEANS_UNCHANGED: Final = frozenset({"duration", "budget_duration"}) + + def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: changed_fields: Final = MappingProxyType( { field: value for field, value in data.model_dump(exclude_unset=True).items() - if field in UpdateKeyRequest.model_fields and field != "key" + if field in UpdateKeyRequest.model_fields + and field != "key" + and not (field in _EMPTY_DURATION_MEANS_UNCHANGED and value == "") } ) if not changed_fields: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 82458b5dfa3..df58424095f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -12142,7 +12142,10 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho @pytest.mark.asyncio -@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +@pytest.mark.parametrize( + "data", + [None, RegenerateKeyRequest(), RegenerateKeyRequest(duration=""), RegenerateKeyRequest(budget_duration="")], +) async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data): mock_prisma_client = _make_regenerate_mock_prisma() @@ -12184,6 +12187,58 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_hides_the_untouched_modal_expiry_from_the_custom_key_update_hook(): + mock_prisma_client = _make_regenerate_mock_prisma() + untouched_modal_body = RegenerateKeyRequest( + key_alias=None, max_budget=None, tpm_limit=None, rpm_limit=None, duration="", grace_period="" + ) + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration is not None and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for the untouched modal body + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=untouched_modal_body, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + assert len(received_data) == 1 + assert "duration" not in received_data[0].model_fields_set + assert received_data[0].model_fields_set >= {"key", "key_alias", "max_budget", "tpm_limit", "rpm_limit"} + + _POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" _POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" _POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} From a635d7be6a10d5126790cd10df7b67fcfbf1a2a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:29:09 -0700 Subject: [PATCH 059/297] fix(guardrails): write per-message guardrail rewrites back onto Responses input items A guardrail that answers one rewritten text per message it saw no longer matches the texts the Responses handler extracted once the request carries instructions or tool items, so the rewrite was rejected with a 500. Spread such an answer over the structured messages' text slots and write it back through the structured path, have Prompt Security modify return structured_messages directly, and give the chat completions pairing the same named rejection instead of a silent misalignment when the counts differ. --- .../base_llm/guardrail_translation/utils.py | 77 ++++++++++++- .../chat/guardrail_translation/handler.py | 4 + .../guardrail_translation/handler.py | 27 ++++- .../prompt_security/prompt_security.py | 39 ++++++- .../test_openai_guardrail_handler.py | 43 ++++++++ ...test_openai_responses_guardrail_handler.py | 103 ++++++++++++++++++ .../test_prompt_security_guardrails.py | 89 +++++++++++++++ 7 files changed, 373 insertions(+), 9 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 94a780f8148..383e668e45c 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,8 +1,10 @@ from __future__ import annotations import json -from collections.abc import Callable, Iterator, Sequence -from typing import Final, TypeVar +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import accumulate +from types import MappingProxyType +from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles from pydantic import BaseModel @@ -364,3 +366,74 @@ def merge_guardrailed_scoped_messages( yield from appended return list(_merged()) + + +def _content_part_text(part: object) -> str | None: + if not isinstance(part, Mapping): + return None + text: Final = part.get("text") + return text if isinstance(text, str) else None + + +def message_text_slot_count(message: AllMessageValues) -> int: + content: Final = message.get("content") + if isinstance(content, str): + return 1 + if isinstance(content, list): + return sum(1 for part in content if _content_part_text(part) is not None) + return 0 + + +def _part_with_text(part: object, text: str) -> object: + if not isinstance(part, Mapping): + return part + return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts + + +def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]: + text_part_indices: Final = tuple( + index for index, part in enumerate(content) if _content_part_text(part) is not None + ) + replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts))) + return [ # mutable-ok: message content stays a JSON list + _part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part + for index, part in enumerate(content) + ] + + +def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues: + content: Final = message.get("content") + if not isinstance(content, (str, list)): + return message + rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) + rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts + return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped + + +def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: + if message_text_slot_count(message) != len(texts): + return None + return _message_with_slot_texts(message, texts) + + +def messages_with_slot_texts( + messages: Sequence[AllMessageValues], + texts: Sequence[str], +) -> list[AllMessageValues] | None: + """Spread one flat list of rewritten texts over the messages' text slots, in order. + + A slot is a string ``content`` or one list part carrying a string ``text``; + images and other parts ride along untouched. A guardrail that answers one + text per message it saw produces exactly this shape, which stops matching + the endpoint's own per-text extraction as soon as the request carries + instructions or tool items. Returns None unless the counts line up exactly, + so a rewrite never lands on the wrong slot. + """ + slot_counts: Final = tuple(message_text_slot_count(message) for message in messages) + if sum(slot_counts) != len(texts): + return None + offsets: Final = tuple(accumulate(slot_counts, initial=0)) + return [ # mutable-ok: guardrail rows travel as a list + _message_with_slot_texts(message, texts[start:end]) + for message, start, end in zip(messages, offsets, offsets[1:]) + ] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 58ff03e6a0d..56fda636e9a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -196,6 +196,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: + if len(guardrailed_texts) != len(text_task_mappings): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") await self._apply_guardrail_responses_to_input_texts( messages=messages, responses=guardrailed_texts, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..61903a54cd3 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -53,6 +53,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, + messages_with_slot_texts, stream_item_field, stream_item_fingerprint, stream_item_items, @@ -395,6 +396,20 @@ def _patched_request_fields( ) +def _guardrailed_structured_messages( + structured_messages: Sequence[AllMessageValues] | None, + sent_text_count: int, + guardrailed_inputs: GenericGuardrailAPIInputs, +) -> Sequence[AllMessageValues] | None: + returned: Final = guardrailed_inputs.get("structured_messages") + if returned is not None and returned is not structured_messages: + return returned + rewritten_texts: Final = guardrailed_inputs.get("texts") + if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count: + return None + return messages_with_slot_texts(structured_messages, rewritten_texts) + + def _patch_or_convert_request_fields( raw_input: object, instructions: object, @@ -473,7 +488,8 @@ class OpenAIResponsesHandler(BaseTranslation): form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) - if not extracted.inputs.get("texts"): + sent_texts: Final = extracted.inputs.get("texts") + if not sent_texts: return data if structured_messages: extracted.inputs["structured_messages"] = structured_messages @@ -486,7 +502,9 @@ class OpenAIResponsesHandler(BaseTranslation): self._apply_guardrailed_tools_to_data( data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) - written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) + written_back: Final = self._written_back_request_fields( + data, structured_messages, len(sent_texts), guardrailed_inputs + ) if written_back is not None: data["input"] = list(written_back.input) # mutable-ok: JSON body if written_back.instructions is None: @@ -553,10 +571,11 @@ class OpenAIResponsesHandler(BaseTranslation): def _written_back_request_fields( data: Mapping[str, object], structured_messages: Sequence[AllMessageValues] | None, + sent_text_count: int, guardrailed_inputs: GenericGuardrailAPIInputs, ) -> _RequestFields | None: - guardrailed: Final = guardrailed_inputs.get("structured_messages") - if guardrailed is None or guardrailed is structured_messages: + guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs) + if guardrailed is None: return None return _patch_or_convert_request_fields( data.get("input"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 0954fe1698a..72c87a793a5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -2,6 +2,7 @@ import asyncio import base64 import os from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional import httpx @@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -27,6 +30,7 @@ if TYPE_CHECKING: _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0 +_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"}) class PromptSecurityGuardrailMissingSecrets(Exception): @@ -275,14 +279,44 @@ class PromptSecurityGuardrail(CustomGuardrail): detail="Blocked by Prompt Security, Violations: " + ", ".join(violations), ) elif action == "modify": - # Extract modified texts from modified_messages modified_messages: Final = result.get("modified_messages", []) modified_texts: Final = self._extract_texts_from_messages(modified_messages) if modified_texts: inputs["texts"] = modified_texts + rewritten_messages: Final = self._structured_messages_with_modifications( + structured_messages, modified_messages + ) + if rewritten_messages is not None: + inputs["structured_messages"] = rewritten_messages return inputs + def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool: + return self.check_tool_results or message.get("role") in _PROTECT_ROLES + + def _structured_messages_with_modifications( + self, + structured_messages: Sequence[AllMessageValues], + modified_messages: Sequence[Mapping[str, object]], + ) -> list[AllMessageValues] | None: + sent_indices: Final = tuple( + index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message) + ) + if not sent_indices or len(sent_indices) != len(modified_messages): + return None + rewritten: Final = tuple( + message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,))) + for index, modified in zip(sent_indices, modified_messages) + ) + replacements: Final = MappingProxyType( + {index: message for index, message in zip(sent_indices, rewritten) if message is not None} + ) + if len(replacements) != len(sent_indices): + return None + return [ # mutable-ok: guardrail inputs take a list + replacements.get(index, message) for index, message in enumerate(structured_messages) + ] + async def _apply_guardrail_on_response( self, inputs: GenericGuardrailAPIInputs, @@ -678,14 +712,13 @@ class PromptSecurityGuardrail(CustomGuardrail): This allows checking tool results for indirect prompt injection when enabled. """ - supported_roles: Final = ["system", "user", "assistant"] filtered_messages: Final = [] transformed_count = 0 filtered_count = 0 for message in messages: role = message.get("role", "") - if role in supported_roles: + if role in _PROTECT_ROLES: filtered_messages.append(message) else: if self.check_tool_results: 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 5a29a96829f..4e2291fdec1 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 @@ -1893,6 +1893,49 @@ class TestScanOnlyToolResults: assert data["messages"][4]["content"] == "and then?" +class ToolDroppingTextGuardrail(CustomGuardrail): + """Answers one text per non-tool message it saw, the way a guardrail that + filters tool rows out before scanning does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="tool-dropping-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"] + return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]} + + +class TestPerMessageTextWriteBack: + """Texts that no longer pair one-to-one with what the handler extracted must be + rejected by name instead of sliding onto the wrong messages.""" + + @pytest.mark.asyncio + async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + handler = OpenAIChatCompletionsHandler() + original_messages = [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + {"role": "user", "content": "fetch the page"}, + {"role": "assistant", "content": "fetching"}, + {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"}, + {"role": "user", "content": "and then?"}, + ] + data = {"messages": json.loads(json.dumps(original_messages))} + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail()) + + assert excinfo.value.guardrail_name == "tool-dropping-redactor" + assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" 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 a4f0a77a9b6..b4b467773da 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 @@ -2338,6 +2338,109 @@ def _parallel_tool_call_input() -> list: ] +SSN = "123-45-6789" +REDACTED_SSN = "" + + +def _slot_texts(message: dict) -> list[str]: + content = message.get("content") + if isinstance(content, str): + return [content] + if isinstance(content, list): + return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)] + return [] + + +class PerMessageRedactionGuardrail(CustomGuardrail): + """Guardrail that answers one redacted text per message it was shown and hands + back only texts, the way Prompt Security in modify mode and a generic guardrail + API server that scans per message do.""" + + def __init__(self, extra_texts: int = 0): + super().__init__(guardrail_name="per-message-redactor") + self.extra_texts = extra_texts + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = inputs.get("structured_messages") or [] + texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)] + return {**inputs, "texts": texts + ["junk"] * self.extra_texts} + + +class TestPerMessageTextWriteBack: + """A guardrail that rewrites one text per message it saw must land on the + instructions and the input items those messages came from, not be rejected.""" + + @pytest.mark.asyncio + async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self): + handler = OpenAIResponsesHandler() + function_call_item = { + "type": "function_call", + "call_id": "call_1", + "name": "lookup_customer", + "arguments": '{"query": "' + SSN + '"}', + } + data = { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + function_call_item, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ], + } + + result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + + assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." + assert [item.get("type", item.get("role")) for item in result["input"]] == [ + "user", + "function_call", + "function_call_output", + ] + assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] + assert result["input"][1] == function_call_item + assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}' + assert result["input"][2]["call_id"] == "call_1" + + @pytest.mark.asyncio + async def test_string_input_with_instructions_keeps_the_two_apart(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "instructions": "Redact " + SSN + " everywhere.", + "input": "My SSN is " + SSN + ".", + } + + result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + + assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere." + assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] + + @pytest.mark.asyncio + async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + handler = OpenAIResponsesHandler() + original_input = [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ] + data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)} + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1)) + + assert excinfo.value.guardrail_name == "per-message-redactor" + assert data["input"] == original_input + assert data["instructions"] == "Be terse." + + class TestProvenancePatching: """The O(n) provenance pass must keep patching rewritten rows in place for the shapes real agent loops produce, and fall back safely everywhere else.""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index e650f796f29..9e83098eb04 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -174,6 +174,95 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] +def _modify_response(modified_messages: list) -> Response: + mock_response = Response( + json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/protect"), + ) + mock_response.raise_for_status = lambda: None + return mock_response + + +def _tool_replay_messages() -> list: + return [ + {"role": "system", "content": "Never echo an SSN like 123-45-6789."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Look up 123-45-6789"}, + {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}}, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + {"role": "user", "content": "Summarize what you found."}, + ] + + +@pytest.mark.asyncio +async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatch: pytest.MonkeyPatch): + """A per-message modify verdict comes back as structured_messages so the + endpoint handler can write it back by message, with the rows Prompt Security + never saw (tool results) and the non-text parts (images) left in place.""" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages = _tool_replay_messages() + inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages} + modified_messages = [ + {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, + {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}]}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "Summarize what you found."}, + ] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + + assert result["structured_messages"] == [ + {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Look up [REDACTED]"}, + {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}}, + ], + }, + messages[2], + messages[3], + {"role": "user", "content": "Summarize what you found."}, + ] + assert result["structured_messages"] is not messages + assert result["texts"] == [ + "Never echo an SSN like [REDACTED].", + "Look up [REDACTED]", + "Summarize what you found.", + ] + + +@pytest.mark.asyncio +async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages = _tool_replay_messages() + inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages} + modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + + assert result["structured_messages"] is messages + assert result["texts"] == ["Look up [REDACTED]"] + + @pytest.mark.asyncio async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail allows safe prompts""" From c1d0d29d01dbd4c28e5a80a52b5af99a73383dd0 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sun, 13 Sep 2026 04:50:29 +0000 Subject: [PATCH 060/297] fix(guardrails): import ModelResponse lazily to avoid cyclic import alert Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 34b79e295e1..1d00ad8c29a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -31,7 +31,6 @@ from litellm.types.utils import ( GuardrailStatus, GuardrailTracingDetail, LLMResponseTypes, - ModelResponse, StandardLoggingGuardrailInformation, ) @@ -907,6 +906,8 @@ class CustomGuardrail(CustomLogger): response: Final = ( kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result ) + from litellm.types.utils import ModelResponse + output_translation: Final = ( get_guardrail_translation_mapping(CallTypes.acompletion)() if isinstance(response, ModelResponse) From 6423acc11a1577d70fa7a0aadd7ece09ca21505b Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 04:55:39 +0000 Subject: [PATCH 061/297] chore(prices): sync prices for 5 providers: 278 models, 34 new fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/deepseek-v4-flash-0731: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp: fireworks_ai/deepseek-v4-flash-vision-exp: fireworks_ai/accounts/fireworks/models/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/deepseek-v4p1-flash: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/glm-5p2: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/glm-5p2: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/glm-5p3: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/glm-5p3-flash: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/gpt-oss-120b: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/gpt-oss-120b: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/kimi-k2p6: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/kimi-k2p6: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/kimi-k2p7-code: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/kimi-k2p7-code: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/kimi-k3: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/kimi-k3: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/minimax-m2p7: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/minimax-m2p7: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/minimax-m3: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/minimax-m3: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/models/muse-glimmer-30b: fireworks_ai/muse-glimmer-30b: fireworks_ai/accounts/fireworks/models/nemotron-3-ultra-nvfp4: fireworks_ai/nemotron-3-ultra-nvfp4: fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b: fireworks_ai/accounts/fireworks/models/qwen3-reranker-8b: input_cost_per_token fireworks_ai/accounts/fireworks/models/qwen3p7-plus: fireworks_ai/qwen3p7-plus: fireworks_ai/accounts/fireworks/models/qwen3p8-max: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/qwen3p8-max: input_cost_per_token_priority, output_cost_per_token_priority, cache_read_input_token_cost_priority fireworks_ai/accounts/fireworks/routers/glm-5p2-fast: fireworks_ai/accounts/fireworks/routers/glm-5p3-fast: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost fireworks_ai/accounts/fireworks/routers/kimi-k3-fast: together_ai/arcee-ai/trinity-mini: input_cost_per_token, output_cost_per_token together_ai/arize-ai/qwen-2-1.5b-instruct: babbage-002: input_cost_per_token_batches, output_cost_per_token_batches chat-latest: chatgpt-image-latest: output_cost_per_token, input_cost_per_image_token, output_cost_per_image_token, input_cost_per_token_batches, output_cost_per_token_batches claude-fable-5: claude-fable-5-1: claude-haiku-4-5: claude-mythos-5: claude-mythos-5-1: claude-opus-4-5: claude-opus-4-6: claude-opus-4-7: claude-opus-4-8: claude-opus-5: claude-sonnet-4-5: claude-sonnet-4-6: claude-sonnet-5: davinci-002: input_cost_per_token_batches, output_cost_per_token_batches deep-research-pro-preview-12-2025: cache_read_input_token_cost together_ai/deepseek-ai/deepseek-coder-33b-instruct: input_cost_per_token, output_cost_per_token together_ai/deepseek-ai/DeepSeek-R1-0528: --- ...odel_prices_and_context_window_backup.json | 968 ++++++++++++++---- model_prices_and_context_window.json | 968 ++++++++++++++---- 2 files changed, 1576 insertions(+), 360 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2220d0e1fe5..879cf894152 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11101,12 +11101,15 @@ "babbage-002": { "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, + "input_cost_per_token_batches": 2e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 4e-07 + "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "bedrock/*/1-month-commitment/cohere.command-light-text-v14": { "input_cost_per_second": 0.001902, @@ -13171,7 +13174,9 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "input_cost_per_second": 0.0001, + "source": "https://developers.openai.com/api/docs/pricing" }, "claude-haiku-4-5-20251001": { "deprecation_date": "2026-10-15", @@ -13219,7 +13224,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-3-7-sonnet-20250219": { "cache_creation_input_token_cost": 3.75e-06, @@ -13378,7 +13384,8 @@ "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-5-20250929": { "deprecation_date": "2026-09-29", @@ -13452,7 +13459,7 @@ }, "supports_output_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-6": { "deprecation_date": "2027-02-17", @@ -13488,7 +13495,8 @@ "prompt_cache_min_tokens": 1024, "provider_specific_entry": { "us": 1.1 - } + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -13665,7 +13673,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6": { "deprecation_date": "2027-02-05", @@ -13702,7 +13711,8 @@ "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_speed": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6-20260205": { "deprecation_date": "2027-02-05", @@ -13777,7 +13787,8 @@ }, "supports_output_config": true, "supports_speed": true, - "prompt_cache_min_tokens": 2048 + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-7-20260416": { "deprecation_date": "2027-04-16", @@ -13855,7 +13866,7 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512, "supports_native_structured_output": true, - "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-5-1": { "deprecation_date": "2027-09-01", @@ -13896,7 +13907,7 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512, "supports_native_structured_output": true, - "source": "https://platform.claude.com/docs/en/models/fable-5-1/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-5": { "deprecation_date": "2027-07-24", @@ -13937,7 +13948,7 @@ "supports_output_config": true, "supports_speed": true, "prompt_cache_min_tokens": 512, - "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-8": { "deprecation_date": "2027-05-28", @@ -13977,7 +13988,8 @@ }, "supports_output_config": true, "supports_speed": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-06-15", @@ -19246,12 +19258,15 @@ "davinci-002": { "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 2e-06 + "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, + "source": "https://developers.openai.com/api/docs/pricing" }, "deepgram/base": { "input_cost_per_second": 0.00020833, @@ -22298,15 +22313,18 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro": { - "cache_read_input_token_cost": 1.45e-07, - "input_cost_per_token": 1.74e-06, + "cache_read_input_token_cost": 6e-07, + "cache_read_input_token_cost_priority": 6e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.48e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22315,14 +22333,17 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost_priority": 5.5e-08, "input_cost_per_token": 1.32e-06, + "input_cost_per_token_priority": 1.65e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 3.96e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 4.95e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22417,14 +22438,17 @@ }, "fireworks_ai/accounts/fireworks/models/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, + "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22433,14 +22457,17 @@ }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "cache_read_input_token_cost": 1.5e-08, + "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.8e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 7.2e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22519,14 +22546,17 @@ }, "fireworks_ai/accounts/fireworks/models/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, + "cache_read_input_token_cost_priority": 2.2e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22535,14 +22565,17 @@ }, "fireworks_ai/accounts/fireworks/models/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, + "cache_read_input_token_cost_priority": 2.85e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22668,14 +22701,17 @@ }, "fireworks_ai/accounts/fireworks/models/minimax-m2p7": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 6e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 196608, "max_output_tokens": 196608, "max_tokens": 196608, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22684,14 +22720,17 @@ }, "fireworks_ai/accounts/fireworks/models/minimax-m3": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 512000, "max_output_tokens": 512000, "max_tokens": 512000, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.8e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22767,15 +22806,18 @@ "supports_vision": false }, "fireworks_ai/deepseek-v4-pro": { - "cache_read_input_token_cost": 1.45e-07, - "input_cost_per_token": 1.74e-06, + "cache_read_input_token_cost": 6e-07, + "cache_read_input_token_cost_priority": 6e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.48e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22831,14 +22873,17 @@ }, "fireworks_ai/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, + "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22847,14 +22892,17 @@ }, "fireworks_ai/gpt-oss-120b": { "cache_read_input_token_cost": 1.5e-08, + "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.8e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 7.2e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22893,14 +22941,17 @@ }, "fireworks_ai/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, + "cache_read_input_token_cost_priority": 2.2e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22925,14 +22976,17 @@ }, "fireworks_ai/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, + "cache_read_input_token_cost_priority": 2.85e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22971,14 +23025,17 @@ }, "fireworks_ai/minimax-m2p7": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 6e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 196608, "max_output_tokens": 196608, "max_tokens": 196608, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22987,14 +23044,17 @@ }, "fireworks_ai/minimax-m3": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 512000, "max_output_tokens": 512000, "max_tokens": 512000, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.8e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -23010,7 +23070,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1.6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -23287,26 +23347,28 @@ "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, - "input_cost_per_token_batches": 2e-07, + "input_cost_per_token_batches": 8e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 1.6e-06, - "output_cost_per_token_batches": 2e-07 + "output_cost_per_token_batches": 9e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "ft:davinci-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.2e-05, - "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_batches": 6e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 1.2e-05, - "output_cost_per_token_batches": 1e-06 + "output_cost_per_token_batches": 6e-06, + "source": "https://developers.openai.com/api/docs/pricing" }, "ft:gpt-3.5-turbo": { "deprecation_date": "2026-10-23", @@ -23319,6 +23381,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_batches": 3e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_system_messages": true, "supports_tool_choice": true }, @@ -23375,14 +23438,15 @@ "ft:gpt-4o-2024-08-06": { "cache_read_input_token_cost": 1.875e-06, "input_cost_per_token": 3.75e-06, - "input_cost_per_token_batches": 1.875e-06, + "input_cost_per_token_batches": 2.225e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-05, - "output_cost_per_token_batches": 7.5e-06, + "output_cost_per_token_batches": 1.25e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -23420,6 +23484,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "output_cost_per_token_batches": 6e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -23439,6 +23504,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -23457,6 +23523,7 @@ "mode": "chat", "output_cost_per_token": 3.2e-06, "output_cost_per_token_batches": 1.6e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -23476,6 +23543,7 @@ "mode": "chat", "output_cost_per_token": 8e-07, "output_cost_per_token_batches": 4e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -23495,6 +23563,7 @@ "mode": "chat", "output_cost_per_token": 1.6e-05, "output_cost_per_token_batches": 8e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_prompt_caching": true, @@ -23505,15 +23574,18 @@ "gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, "deprecation_date": "2026-06-01", - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_character": 3.75e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", + "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text", "image", @@ -23582,13 +23654,16 @@ "cache_read_input_token_cost": 1.875e-08, "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, + "input_cost_per_character": 1.875e-08, "input_cost_per_token": 7.5e-08, + "input_cost_per_token_batches": 3.75e-08, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 3e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_token_batches": 1.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text", "image", @@ -23651,6 +23726,8 @@ "gemini-2.5-flash": { "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 3e-08, + "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", @@ -23660,7 +23737,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23693,6 +23770,12 @@ "search_context_size_high": 0.035 }, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, "supports_image_size": false }, "gemini-2.5-flash-image": { @@ -23700,6 +23783,9 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, @@ -23709,8 +23795,10 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23741,10 +23829,19 @@ "supports_image_size": false }, "gemini-3-pro-image": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -23753,8 +23850,12 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23822,9 +23923,13 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -23833,7 +23938,9 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23900,9 +24007,11 @@ }, "gemini-3.1-flash-lite-image": { "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_flex": 1.25e-08, "input_cost_per_image": 0.00028, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 4096, @@ -23912,6 +24021,7 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -24005,7 +24115,7 @@ "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, "output_cost_per_token_priority": 2.7e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24047,7 +24157,7 @@ "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 1.5e-08, - "cache_read_input_token_cost_priority": 5e-08, + "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, @@ -24062,7 +24172,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24101,6 +24211,7 @@ "google_maps_grounding_cost_per_query": 0.014 }, "deep-research-pro-preview-12-2025": { + "cache_read_input_token_cost": 2e-07, "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -24113,7 +24224,7 @@ "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24137,6 +24248,8 @@ "gemini-2.5-flash-lite": { "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_flex": 1e-08, + "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -24146,7 +24259,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 4e-07, "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24179,6 +24292,12 @@ "search_context_size_high": 0.035 }, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_batches": 5e-08, + "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.8e-07, + "output_cost_per_token_batches": 2e-07, + "output_cost_per_token_flex": 2e-07, + "output_cost_per_token_priority": 7.2e-07, "supports_image_size": false }, "gemini-2.5-flash-lite-preview-09-2025": { @@ -24330,7 +24449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/vertex_ai/live" ], @@ -24361,7 +24480,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "gemini_native_audio": true + "gemini_native_audio": true, + "input_cost_per_image_token": 3e-06 }, "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, @@ -24461,6 +24581,9 @@ "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, + "cache_read_input_token_cost_flex": 1.25e-07, + "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", @@ -24470,7 +24593,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions" @@ -24500,7 +24623,15 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_flex": 6.25e-07, + "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_flex": 5e-06, + "output_cost_per_token_priority": 1.8e-05 }, "gemini-3-pro-preview": { "deprecation_date": "2026-03-26", @@ -24575,7 +24706,7 @@ "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_image": 0.00012, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24609,13 +24740,16 @@ "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_flex": 2e-07, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -25127,13 +25261,15 @@ "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", + "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2e-05, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -25340,7 +25476,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text", "image" @@ -25381,8 +25517,11 @@ }, "gemini-embedding-2": { "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image": 0.00012, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_per_second": 0.00079, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, @@ -25390,7 +25529,7 @@ "mode": "embedding", "output_cost_per_token": 0, "output_vector_size": 3072, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_multimodal": true, "uses_embed_content": true }, @@ -27055,6 +27194,7 @@ }, "gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -27064,7 +27204,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27102,7 +27242,11 @@ "search_context_size_high": 0.014 }, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06 }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -27115,7 +27259,7 @@ "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "output_cost_per_video_token": 1.75e-05, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions" ], @@ -27148,7 +27292,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27211,7 +27355,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27268,7 +27412,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27325,7 +27469,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -28783,6 +28927,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_system_messages": true, @@ -28791,12 +28936,15 @@ "gpt-3.5-turbo-0125": { "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -28806,12 +28954,15 @@ "gpt-3.5-turbo-1106": { "deprecation_date": "2026-09-28", "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 2e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -28839,7 +28990,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 2e-06 + "output_cost_per_token": 2e-06, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-3.5-turbo-instruct-0914": { "input_cost_per_token": 1.5e-06, @@ -28894,12 +29046,15 @@ "gpt-4-0613": { "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 8192, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-05, + "output_cost_per_token_batches": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_system_messages": true, @@ -28940,12 +29095,15 @@ "gpt-4-turbo-2024-04-09": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3e-05, + "output_cost_per_token_batches": 1.5e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -28989,6 +29147,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29031,6 +29190,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29073,6 +29233,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29115,6 +29276,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29153,6 +29315,7 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29190,6 +29353,7 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_priority": 8e-07, "output_cost_per_token_batches": 2e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29226,6 +29390,7 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29248,6 +29413,7 @@ "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 2.625e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29270,6 +29436,7 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29293,6 +29460,7 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29367,6 +29535,7 @@ "mode": "chat", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29403,6 +29572,7 @@ "mode": "chat", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions" ], @@ -29437,6 +29607,7 @@ "mode": "chat", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29474,6 +29645,7 @@ "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29511,6 +29683,7 @@ "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29572,7 +29745,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": false, - "deprecation_date": "2027-01-20" + "deprecation_date": "2027-01-20", + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -29600,7 +29774,8 @@ "search_context_size_high": 0.025, "search_context_size_low": 0.025, "search_context_size_medium": 0.025 - } + }, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29621,6 +29796,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29769,15 +29945,18 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "input_cost_per_second": 5e-05, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-tts": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_second": 0.00025, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -29909,7 +30088,9 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "input_cost_per_second": 0.0001, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-image-1.5": { "cache_read_input_token_cost": 1.25e-06, @@ -29919,7 +30100,10 @@ "mode": "image_generation", "output_cost_per_token": 1e-05, "input_cost_per_image_token": 8e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations" ], @@ -29934,7 +30118,10 @@ "mode": "image_generation", "output_cost_per_token": 1e-05, "input_cost_per_image_token": 8e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations" ], @@ -29947,7 +30134,9 @@ "litellm_provider": "openai", "mode": "image_generation", "input_cost_per_image_token": 8e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -30394,6 +30583,7 @@ "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -30402,6 +30592,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "search_context_cost_per_query": { @@ -30409,6 +30600,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -30438,6 +30630,7 @@ }, "gpt-5.1": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, @@ -30478,11 +30671,17 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_flex": 6.25e-07, + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_flex": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": false, "supports_minimal_reasoning_effort": false }, "gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, @@ -30523,6 +30722,11 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_flex": 6.25e-07, + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_flex": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": false, "supports_minimal_reasoning_effort": false }, @@ -30574,6 +30778,7 @@ }, "gpt-5.2": { "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_flex": 8.75e-08, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, @@ -30615,11 +30820,17 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 8.75e-07, + "input_cost_per_token_flex": 8.75e-07, + "output_cost_per_token_batches": 7e-06, + "output_cost_per_token_flex": 7e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, "gpt-5.2-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_flex": 8.75e-08, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, @@ -30661,6 +30872,11 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 8.75e-07, + "input_cost_per_token_flex": 8.75e-07, + "output_cost_per_token_batches": 7e-06, + "output_cost_per_token_flex": 7e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -30754,17 +30970,20 @@ }, "gpt-5.2-pro": { "input_cost_per_token": 2.1e-05, + "input_cost_per_token_batches": 1.05e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.000168, + "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -30793,17 +31012,20 @@ }, "gpt-5.2-pro-2025-12-11": { "input_cost_per_token": 2.1e-05, + "input_cost_per_token_batches": 1.05e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.000168, + "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -30869,6 +31091,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31005,6 +31228,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31073,6 +31297,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31140,6 +31365,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31203,7 +31429,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://developers.openai.com/api/docs/models/gpt-5.6-cyber", + "source": "https://developers.openai.com/api/docs/pricing", "supports_computer_use": true, "supports_parallel_function_calling": true }, @@ -31373,7 +31599,7 @@ "reasoning_effort_levels": [ "medium" ], - "source": "https://developers.openai.com/api/docs/models/chat-latest", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -31452,7 +31678,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, @@ -31509,7 +31736,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.5-pro": { "input_cost_per_token": 3e-05, @@ -31532,6 +31760,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -31580,6 +31809,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -31657,7 +31887,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -31709,7 +31940,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-pro": { "input_cost_per_token": 3e-05, @@ -31758,7 +31990,8 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, - "output_cost_per_token_above_272k_tokens_flex": 0.000135 + "output_cost_per_token_above_272k_tokens_flex": 0.000135, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-pro-2026-03-05": { "input_cost_per_token": 3e-05, @@ -31807,7 +32040,8 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, - "output_cost_per_token_above_272k_tokens_flex": 0.000135 + "output_cost_per_token_above_272k_tokens_flex": 0.000135, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -31858,6 +32092,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -31910,6 +32145,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -31959,6 +32195,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -32008,6 +32245,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -32026,6 +32264,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -32068,6 +32307,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -32100,6 +32340,7 @@ "cache_read_input_token_cost_priority": 2.5e-07, "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -32108,6 +32349,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "search_context_cost_per_query": { @@ -32115,6 +32357,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32439,6 +32682,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses" ], @@ -32469,6 +32713,7 @@ "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", @@ -32477,6 +32722,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, "search_context_cost_per_query": { @@ -32484,6 +32730,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32517,6 +32764,7 @@ "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", @@ -32525,6 +32773,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, "search_context_cost_per_query": { @@ -32532,6 +32781,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32563,6 +32813,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, + "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -32571,12 +32822,14 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32609,6 +32862,7 @@ "cache_read_input_token_cost_flex": 2.5e-09, "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, + "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", @@ -32617,12 +32871,14 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32655,9 +32911,11 @@ "deprecation_date": "2026-10-23", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 4e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -32668,9 +32926,11 @@ "deprecation_date": "2026-12-01", "input_cost_per_image_token": 2.5e-06, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 8e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -32690,6 +32950,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32722,6 +32983,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32755,6 +33017,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32790,6 +33053,7 @@ "output_cost_per_token": 2.4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32825,6 +33089,7 @@ "output_cost_per_token": 2.4e-06, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32847,8 +33112,10 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, + "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, "litellm_provider": "openai", "max_input_tokens": 32000, @@ -32857,6 +33124,7 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32890,6 +33158,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -37609,12 +37878,15 @@ "cache_read_input_token_cost": 7.5e-06, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 6e-05, + "output_cost_per_token_batches": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, @@ -37629,12 +37901,15 @@ "cache_read_input_token_cost": 7.5e-06, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 6e-05, + "output_cost_per_token_batches": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -37656,6 +37931,7 @@ "mode": "responses", "output_cost_per_token": 0.0006, "output_cost_per_token_batches": 0.0003, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37689,6 +37965,7 @@ "mode": "responses", "output_cost_per_token": 0.0006, "output_cost_per_token_batches": 0.0003, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37716,6 +37993,7 @@ "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -37724,6 +38002,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, "search_context_cost_per_query": { @@ -37731,6 +38010,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/chat/completions", @@ -37760,6 +38040,7 @@ "cache_read_input_token_cost_priority": 8.75e-07, "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -37768,6 +38049,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, "search_context_cost_per_query": { @@ -37775,6 +38057,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/chat/completions", @@ -37884,12 +38167,15 @@ "cache_read_input_token_cost": 5.5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_prompt_caching": true, @@ -37902,12 +38188,15 @@ "cache_read_input_token_cost": 5.5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_prompt_caching": true, @@ -37931,6 +38220,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37968,6 +38258,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37991,10 +38282,11 @@ }, "o4-mini": { "cache_read_input_token_cost": 2.75e-07, - "cache_read_input_token_cost_flex": 1.375e-07, + "cache_read_input_token_cost_flex": 1.38e-07, "cache_read_input_token_cost_priority": 5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, "litellm_provider": "openai", @@ -38003,6 +38295,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "output_cost_per_token_flex": 2.2e-06, "output_cost_per_token_priority": 8e-06, "search_context_cost_per_query": { @@ -38010,6 +38303,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, @@ -38022,10 +38316,11 @@ }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, - "cache_read_input_token_cost_flex": 1.375e-07, + "cache_read_input_token_cost_flex": 1.38e-07, "cache_read_input_token_cost_priority": 5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, "litellm_provider": "openai", @@ -38034,6 +38329,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "output_cost_per_token_flex": 2.2e-06, "output_cost_per_token_priority": 8e-06, "search_context_cost_per_query": { @@ -38041,6 +38337,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, @@ -43203,7 +43500,8 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_cost_per_token_batches": 0.0, - "output_vector_size": 3072 + "output_vector_size": 3072, + "source": "https://developers.openai.com/api/docs/pricing" }, "text-embedding-3-small": { "input_cost_per_token": 2e-08, @@ -43214,7 +43512,8 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_cost_per_token_batches": 0.0, - "output_vector_size": 1536 + "output_vector_size": 1536, + "source": "https://developers.openai.com/api/docs/pricing" }, "text-embedding-ada-002": { "input_cost_per_token": 1e-07, @@ -43223,7 +43522,8 @@ "max_tokens": 8191, "mode": "embedding", "output_cost_per_token": 0.0, - "output_vector_size": 1536 + "output_vector_size": 1536, + "source": "https://developers.openai.com/api/docs/pricing" }, "text-embedding-ada-002-v2": { "input_cost_per_token": 1e-07, @@ -43394,7 +43694,7 @@ "input_cost_per_token": 1.2e-06, "output_cost_per_token": 1.2e-06, "max_input_tokens": 131072, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo": { "litellm_provider": "together_ai", @@ -43406,7 +43706,7 @@ "input_cost_per_token": 3e-07, "output_cost_per_token": 3e-07, "max_input_tokens": 32768, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { "deprecation_date": "2026-07-10", @@ -43453,7 +43753,7 @@ "max_input_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://www.together.ai/models/qwen3-coder-480b-a35b-instruct", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43515,7 +43815,7 @@ }, "mode": "chat", "output_cost_per_token": 1.7e-06, - "source": "https://www.together.ai/models/deepseek-v3-1", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43539,7 +43839,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.04e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43573,6 +43873,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 5.9e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43595,6 +43896,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 8.8e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43606,6 +43908,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 1.8e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43622,7 +43925,7 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, "max_input_tokens": 32768, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { "deprecation_date": "2026-04-02", @@ -43634,7 +43937,7 @@ "input_cost_per_token": 1e-07, "output_cost_per_token": 3e-07, "max_input_tokens": 32768, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { "deprecation_date": "2026-04-16", @@ -43642,6 +43945,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 6e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43668,7 +43972,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://www.together.ai/models/gpt-oss-120b", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43682,7 +43986,7 @@ "max_input_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-07, - "source": "https://www.together.ai/models/gpt-oss-20b", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43702,7 +44006,7 @@ "max_input_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.1e-06, - "source": "https://www.together.ai/models/glm-4-5-air", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43718,7 +44022,7 @@ }, "mode": "chat", "output_cost_per_token": 2.2e-06, - "source": "https://www.together.ai/models/glm-4-6", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43735,7 +44039,7 @@ }, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://www.together.ai/models/glm-4-7", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43783,7 +44087,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43799,7 +44103,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43813,7 +44117,7 @@ "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 3.6e-06, - "source": "https://www.together.ai/models/qwen3-5-397b-a17b", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43828,7 +44132,7 @@ "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43853,7 +44157,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 2.5e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43868,7 +44172,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 3e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { @@ -43879,7 +44183,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 7.5e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/Qwen/Qwen3.7-Plus": { @@ -43889,7 +44193,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.28e-06, - "source": "https://docs.together.ai/docs/serverless-models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { "cache_read_input_token_cost": 2.5e-07, @@ -43899,7 +44203,7 @@ "max_tokens": 1010000, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/arize-ai/qwen-2-1.5b-instruct": { @@ -43909,7 +44213,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1e-07, - "source": "https://docs.together.ai/docs/serverless-models" + "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 3e-08, @@ -43919,7 +44223,7 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 2.8e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43951,7 +44255,7 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 3.96e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43976,7 +44280,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 9.7e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -44012,7 +44316,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/moonshotai/Kimi-K2.7-Code": { @@ -44024,7 +44328,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44045,7 +44349,7 @@ "high", "max" ], - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44063,7 +44367,7 @@ "max_tokens": 512288, "mode": "chat", "output_cost_per_token": 3.6e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44089,7 +44393,7 @@ "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 4.05e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44105,7 +44409,7 @@ "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/zai-org/GLM-5.2": { @@ -44117,7 +44421,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44134,7 +44438,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44151,7 +44455,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44164,6 +44468,7 @@ "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", "mode": "audio_speech", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ] @@ -44172,6 +44477,7 @@ "input_cost_per_character": 3e-05, "litellm_provider": "openai", "mode": "audio_speech", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ] @@ -49497,7 +49803,8 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "source": "https://developers.openai.com/api/docs/pricing" }, "xai/grok-3": { "cache_read_input_token_cost": 2e-07, @@ -52594,10 +52901,11 @@ "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "input_cost_per_token": 0.0, + "input_cost_per_token": 2e-07, "output_cost_per_token": 0.0, "litellm_provider": "fireworks_ai", - "mode": "rerank" + "mode": "rerank", + "source": "https://api.fireworks.ai/v1/serverless/models" }, "fireworks_ai/accounts/fireworks/models/qwen3-vl-235b-a22b-instruct": { "max_tokens": 262144, @@ -52662,7 +52970,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1.6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -54561,12 +54869,13 @@ }, "gpt-4o-mini-tts-2025-03-20": { "deprecation_date": "2026-07-23", - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_second": 0.00025, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -54579,12 +54888,13 @@ ] }, "gpt-4o-mini-tts-2025-12-15": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_second": 0.00025, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -54599,24 +54909,28 @@ "gpt-4o-mini-transcribe-2025-03-20": { "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", "output_cost_per_token": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions" ] }, "gpt-4o-mini-transcribe-2025-12-15": { "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", "output_cost_per_token": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions" ] @@ -54635,6 +54949,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -54662,6 +54977,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -54689,6 +55005,7 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -54740,13 +55057,14 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "deprecation_date": "2027-01-20" + "deprecation_date": "2027-01-20", + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-realtime-whisper": { - "input_cost_per_second": 0.0002833333333333333, + "input_cost_per_second": 0.000283333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://developers.openai.com/api/docs/models/gpt-realtime-whisper", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -54764,7 +55082,7 @@ "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, - "source": "https://platform.openai.com/docs/api-reference/videos", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "text", "image" @@ -54778,7 +55096,7 @@ "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, - "source": "https://platform.openai.com/docs/api-reference/videos", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "text", "image" @@ -54803,11 +55121,15 @@ "chatgpt-image-latest": { "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2026-12-01", - "input_cost_per_image_token": 1e-05, + "input_cost_per_image_token": 8e-06, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "openai", "mode": "image_generation", - "output_cost_per_image_token": 4e-05, + "output_cost_per_image_token": 3.2e-05, + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -57282,7 +57604,7 @@ "input_cost_per_second": 7.5e-05, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://developers.openai.com/api/docs/models/gpt-transcribe", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions", "/v1/realtime/transcription_sessions" @@ -57297,10 +57619,10 @@ "supports_audio_input": true }, "gpt-live-transcribe": { - "input_cost_per_second": 0.0002833333333333333, + "input_cost_per_second": 0.000283333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://developers.openai.com/api/docs/models/gpt-live-transcribe", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -57315,10 +57637,10 @@ "supports_audio_input": true }, "gpt-live-1": { - "input_cost_per_second": 0.0008333333333333334, + "input_cost_per_second": 0.000833333333333, "litellm_provider": "openai", "mode": "realtime", - "source": "https://developers.openai.com/api/docs/models/gpt-live-1", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "text", "audio" @@ -57332,13 +57654,13 @@ "supports_function_calling": true }, "gpt-realtime-translate": { - "input_cost_per_second": 0.0005666666666666667, + "input_cost_per_second": 0.000566666666667, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "max_tokens": 2000, "mode": "realtime", - "source": "https://developers.openai.com/api/docs/models/gpt-realtime-translate", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "audio" ], @@ -57366,7 +57688,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "source": "https://platform.claude.com/docs/en/about-claude/models/overview", + "source": "https://platform.claude.com/docs/en/about-claude/pricing", "supports_adaptive_thinking": true, "thinking_always_on": true, "supports_mid_conversation_system": true, @@ -57427,7 +57749,7 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512, "supports_native_structured_output": true, - "source": "https://platform.claude.com/docs/en/models/mythos-5-1/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-mythos-preview": { "cache_creation_input_token_cost": 1.25e-05, @@ -57779,6 +58101,7 @@ }, "vertex_ai/gemini-3.5-live-translate-preview": { "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_second": 8.83333333333e-05, "input_cost_per_token": 3.5e-06, "litellm_provider": "vertex_ai", "mode": "realtime", @@ -57881,14 +58204,17 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -57897,14 +58223,17 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -57920,26 +58249,29 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/kimi-k3": { "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_priority": 3.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_priority": 1.875e-05, "reasoning_effort_levels": [ "low", "high", "max" ], - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -57948,14 +58280,17 @@ }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -57964,14 +58299,17 @@ }, "fireworks_ai/deepseek-v4p1-flash": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -57987,7 +58325,7 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true @@ -58026,19 +58364,22 @@ }, "fireworks_ai/kimi-k3": { "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_priority": 3.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_priority": 1.875e-05, "reasoning_effort_levels": [ "low", "high", "max" ], - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58089,12 +58430,15 @@ }, "fireworks_ai/qwen3p8-max": { "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 9e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58110,7 +58454,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58142,7 +58486,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58158,7 +58502,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58190,7 +58534,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58199,12 +58543,15 @@ }, "fireworks_ai/accounts/fireworks/models/qwen3p8-max": { "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 9e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58220,7 +58567,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6.6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58257,7 +58604,7 @@ "high", "max" ], - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -60955,14 +61302,17 @@ }, "fireworks_ai/accounts/fireworks/models/glm-5p3": { "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_priority": 3.25e-07, "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -60971,13 +61321,16 @@ }, "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_priority": 3.75e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.875e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -61005,7 +61358,7 @@ "max_output_tokens": 40960, "max_tokens": 40960, "mode": "embedding", - "source": "https://docs.fireworks.ai/serverless/pricing" + "source": "https://api.fireworks.ai/v1/serverless/models" }, "zai/glm-5.2": { "cache_creation_input_token_cost": 0, @@ -61029,7 +61382,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 4.7e-07, - "source": "https://docs.together.ai/docs/serverless-models" + "source": "https://api.together.ai/v1/models" }, "together_ai/moonshotai/Kimi-K2.6": { "deprecation_date": "2026-08-19", @@ -61039,7 +61392,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/moonshotai/Kimi-K2.5-fp4": { "input_cost_per_token": 5e-07, @@ -61047,7 +61400,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/MiniMaxAI/MiniMax-M2.7": { "input_cost_per_token": 3e-07, @@ -61056,7 +61409,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 196608, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/zai-org/GLM-5": { "deprecation_date": "2026-06-22", @@ -61065,7 +61418,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 202752, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/zai-org/GLM-5.1": { "deprecation_date": "2026-07-10", @@ -61075,7 +61428,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 202752, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-0528": { "input_cost_per_token": 3e-06, @@ -61083,7 +61436,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 163840, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-Coder-Next-FP8": { "deprecation_date": "2026-05-14", @@ -61092,7 +61445,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-VL-32B-Instruct": { "deprecation_date": "2026-02-25", @@ -61101,7 +61454,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-VL-8B-Instruct": { "deprecation_date": "2026-04-16", @@ -61110,7 +61463,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Ministral-3-14B-Instruct-2512": { "input_cost_per_token": 2e-07, @@ -61118,7 +61471,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { "input_cost_per_token": 6e-08, @@ -61126,7 +61479,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 131072, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Mistral-7B-Instruct-v0.3": { "input_cost_per_token": 2e-07, @@ -61134,7 +61487,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 32768, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/QwQ-32B": { "deprecation_date": "2025-11-13", @@ -61143,7 +61496,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 131072, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "cerebras/gemma-4-31b": { "input_cost_per_token": 9.9e-07, @@ -65270,5 +65623,260 @@ "supports_tool_choice": false, "supports_response_schema": true, "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models" + }, + "together_ai/arcee-ai/trinity-mini": { + "input_cost_per_token": 4.5e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/deepseek-coder-33b-instruct": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "input_cost_per_token": 2e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "input_cost_per_token": 1.6e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models" + }, + "vertex_ai/gemini-2.5-flash-native-audio": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-2.5-flash-preview-tts": { + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_audio_token": 1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.1-flash-live-preview": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_second": 8.33333333333e-05, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.1-flash-tts-preview": { + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.5-transcribe": { + "input_cost_per_audio_token": 2e-06, + "input_cost_per_second": 5e-05, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.5-transcribe-live": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_second": 8.33333333333e-05, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-omni-1.1-flash": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 9e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-robotics-er-2": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemma-4-26b-a4b-it": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "together_ai/google/gemma-2-27b-it": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "gpt-5.5-cyber": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 1.25e-05, + "litellm_provider": "openai", + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-rosalind-research": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "together_ai/meta-llama/Llama-3-8b-chat-hf": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Llama-3.1-405B-Instruct": { + "input_cost_per_token": 3.5e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 3.5e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Llama-3.2-1B-Instruct": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 6e-08, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Llama-3.2-3B-Instruct": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 6e-08, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Meta-Llama-3-70B-Instruct-Turbo": { + "input_cost_per_token": 8.8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8.8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Meta-Llama-3-8B-Instruct": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF": { + "input_cost_per_token": 8.8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8.8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2-1.5B-Instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-08, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2-72B-Instruct": { + "input_cost_per_token": 9e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2-VL-72B-Instruct": { + "input_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-14B-Instruct": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-72B-Instruct": { + "input_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-Coder-32B-Instruct": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-VL-72B-Instruct": { + "input_cost_per_token": 1.95e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://api.together.ai/v1/models" } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2220d0e1fe5..879cf894152 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11101,12 +11101,15 @@ "babbage-002": { "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, + "input_cost_per_token_batches": 2e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 4e-07 + "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "bedrock/*/1-month-commitment/cohere.command-light-text-v14": { "input_cost_per_second": 0.001902, @@ -13171,7 +13174,9 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "input_cost_per_second": 0.0001, + "source": "https://developers.openai.com/api/docs/pricing" }, "claude-haiku-4-5-20251001": { "deprecation_date": "2026-10-15", @@ -13219,7 +13224,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-3-7-sonnet-20250219": { "cache_creation_input_token_cost": 3.75e-06, @@ -13378,7 +13384,8 @@ "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-5-20250929": { "deprecation_date": "2026-09-29", @@ -13452,7 +13459,7 @@ }, "supports_output_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-6": { "deprecation_date": "2027-02-17", @@ -13488,7 +13495,8 @@ "prompt_cache_min_tokens": 1024, "provider_specific_entry": { "us": 1.1 - } + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -13665,7 +13673,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6": { "deprecation_date": "2027-02-05", @@ -13702,7 +13711,8 @@ "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_speed": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6-20260205": { "deprecation_date": "2027-02-05", @@ -13777,7 +13787,8 @@ }, "supports_output_config": true, "supports_speed": true, - "prompt_cache_min_tokens": 2048 + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-7-20260416": { "deprecation_date": "2027-04-16", @@ -13855,7 +13866,7 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512, "supports_native_structured_output": true, - "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-5-1": { "deprecation_date": "2027-09-01", @@ -13896,7 +13907,7 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512, "supports_native_structured_output": true, - "source": "https://platform.claude.com/docs/en/models/fable-5-1/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-5": { "deprecation_date": "2027-07-24", @@ -13937,7 +13948,7 @@ "supports_output_config": true, "supports_speed": true, "prompt_cache_min_tokens": 512, - "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-8": { "deprecation_date": "2027-05-28", @@ -13977,7 +13988,8 @@ }, "supports_output_config": true, "supports_speed": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-06-15", @@ -19246,12 +19258,15 @@ "davinci-002": { "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 2e-06 + "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, + "source": "https://developers.openai.com/api/docs/pricing" }, "deepgram/base": { "input_cost_per_second": 0.00020833, @@ -22298,15 +22313,18 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro": { - "cache_read_input_token_cost": 1.45e-07, - "input_cost_per_token": 1.74e-06, + "cache_read_input_token_cost": 6e-07, + "cache_read_input_token_cost_priority": 6e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.48e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22315,14 +22333,17 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost_priority": 5.5e-08, "input_cost_per_token": 1.32e-06, + "input_cost_per_token_priority": 1.65e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 3.96e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 4.95e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22417,14 +22438,17 @@ }, "fireworks_ai/accounts/fireworks/models/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, + "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22433,14 +22457,17 @@ }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "cache_read_input_token_cost": 1.5e-08, + "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.8e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 7.2e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22519,14 +22546,17 @@ }, "fireworks_ai/accounts/fireworks/models/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, + "cache_read_input_token_cost_priority": 2.2e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22535,14 +22565,17 @@ }, "fireworks_ai/accounts/fireworks/models/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, + "cache_read_input_token_cost_priority": 2.85e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22668,14 +22701,17 @@ }, "fireworks_ai/accounts/fireworks/models/minimax-m2p7": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 6e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 196608, "max_output_tokens": 196608, "max_tokens": 196608, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22684,14 +22720,17 @@ }, "fireworks_ai/accounts/fireworks/models/minimax-m3": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 512000, "max_output_tokens": 512000, "max_tokens": 512000, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.8e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22767,15 +22806,18 @@ "supports_vision": false }, "fireworks_ai/deepseek-v4-pro": { - "cache_read_input_token_cost": 1.45e-07, - "input_cost_per_token": 1.74e-06, + "cache_read_input_token_cost": 6e-07, + "cache_read_input_token_cost_priority": 6e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.48e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22831,14 +22873,17 @@ }, "fireworks_ai/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, + "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22847,14 +22892,17 @@ }, "fireworks_ai/gpt-oss-120b": { "cache_read_input_token_cost": 1.5e-08, + "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.8e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 7.2e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22893,14 +22941,17 @@ }, "fireworks_ai/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, + "cache_read_input_token_cost_priority": 2.2e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22925,14 +22976,17 @@ }, "fireworks_ai/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, + "cache_read_input_token_cost_priority": 2.85e-07, "input_cost_per_token": 9.5e-07, + "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22971,14 +23025,17 @@ }, "fireworks_ai/minimax-m2p7": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 6e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 1.2e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 196608, "max_output_tokens": 196608, "max_tokens": 196608, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.2e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -22987,14 +23044,17 @@ }, "fireworks_ai/minimax-m3": { "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 512000, "max_output_tokens": 512000, "max_tokens": 512000, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 1.8e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -23010,7 +23070,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1.6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -23287,26 +23347,28 @@ "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, - "input_cost_per_token_batches": 2e-07, + "input_cost_per_token_batches": 8e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 1.6e-06, - "output_cost_per_token_batches": 2e-07 + "output_cost_per_token_batches": 9e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "ft:davinci-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.2e-05, - "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_batches": 6e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 1.2e-05, - "output_cost_per_token_batches": 1e-06 + "output_cost_per_token_batches": 6e-06, + "source": "https://developers.openai.com/api/docs/pricing" }, "ft:gpt-3.5-turbo": { "deprecation_date": "2026-10-23", @@ -23319,6 +23381,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_batches": 3e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_system_messages": true, "supports_tool_choice": true }, @@ -23375,14 +23438,15 @@ "ft:gpt-4o-2024-08-06": { "cache_read_input_token_cost": 1.875e-06, "input_cost_per_token": 3.75e-06, - "input_cost_per_token_batches": 1.875e-06, + "input_cost_per_token_batches": 2.225e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-05, - "output_cost_per_token_batches": 7.5e-06, + "output_cost_per_token_batches": 1.25e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -23420,6 +23484,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "output_cost_per_token_batches": 6e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -23439,6 +23504,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -23457,6 +23523,7 @@ "mode": "chat", "output_cost_per_token": 3.2e-06, "output_cost_per_token_batches": 1.6e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -23476,6 +23543,7 @@ "mode": "chat", "output_cost_per_token": 8e-07, "output_cost_per_token_batches": 4e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -23495,6 +23563,7 @@ "mode": "chat", "output_cost_per_token": 1.6e-05, "output_cost_per_token_batches": 8e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_prompt_caching": true, @@ -23505,15 +23574,18 @@ "gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, "deprecation_date": "2026-06-01", - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_character": 3.75e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", + "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text", "image", @@ -23582,13 +23654,16 @@ "cache_read_input_token_cost": 1.875e-08, "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, + "input_cost_per_character": 1.875e-08, "input_cost_per_token": 7.5e-08, + "input_cost_per_token_batches": 3.75e-08, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 3e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_token_batches": 1.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text", "image", @@ -23651,6 +23726,8 @@ "gemini-2.5-flash": { "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 3e-08, + "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", @@ -23660,7 +23737,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23693,6 +23770,12 @@ "search_context_size_high": 0.035 }, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, "supports_image_size": false }, "gemini-2.5-flash-image": { @@ -23700,6 +23783,9 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, @@ -23709,8 +23795,10 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23741,10 +23829,19 @@ "supports_image_size": false }, "gemini-3-pro-image": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -23753,8 +23850,12 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23822,9 +23923,13 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -23833,7 +23938,9 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23900,9 +24007,11 @@ }, "gemini-3.1-flash-lite-image": { "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_flex": 1.25e-08, "input_cost_per_image": 0.00028, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 4096, @@ -23912,6 +24021,7 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -24005,7 +24115,7 @@ "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, "output_cost_per_token_priority": 2.7e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24047,7 +24157,7 @@ "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 1.5e-08, - "cache_read_input_token_cost_priority": 5e-08, + "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, @@ -24062,7 +24172,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24101,6 +24211,7 @@ "google_maps_grounding_cost_per_query": 0.014 }, "deep-research-pro-preview-12-2025": { + "cache_read_input_token_cost": 2e-07, "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -24113,7 +24224,7 @@ "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24137,6 +24248,8 @@ "gemini-2.5-flash-lite": { "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_flex": 1e-08, + "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -24146,7 +24259,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 4e-07, "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24179,6 +24292,12 @@ "search_context_size_high": 0.035 }, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_batches": 5e-08, + "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.8e-07, + "output_cost_per_token_batches": 2e-07, + "output_cost_per_token_flex": 2e-07, + "output_cost_per_token_priority": 7.2e-07, "supports_image_size": false }, "gemini-2.5-flash-lite-preview-09-2025": { @@ -24330,7 +24449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/vertex_ai/live" ], @@ -24361,7 +24480,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "gemini_native_audio": true + "gemini_native_audio": true, + "input_cost_per_image_token": 3e-06 }, "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, @@ -24461,6 +24581,9 @@ "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, + "cache_read_input_token_cost_flex": 1.25e-07, + "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", @@ -24470,7 +24593,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions" @@ -24500,7 +24623,15 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_flex": 6.25e-07, + "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_flex": 5e-06, + "output_cost_per_token_priority": 1.8e-05 }, "gemini-3-pro-preview": { "deprecation_date": "2026-03-26", @@ -24575,7 +24706,7 @@ "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_image": 0.00012, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -24609,13 +24740,16 @@ "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_flex": 2e-07, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -25127,13 +25261,15 @@ "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", + "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2e-05, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -25340,7 +25476,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_modalities": [ "text", "image" @@ -25381,8 +25517,11 @@ }, "gemini-embedding-2": { "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image": 0.00012, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_per_second": 0.00079, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, @@ -25390,7 +25529,7 @@ "mode": "embedding", "output_cost_per_token": 0, "output_vector_size": 3072, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_multimodal": true, "uses_embed_content": true }, @@ -27055,6 +27194,7 @@ }, "gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -27064,7 +27204,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27102,7 +27242,11 @@ "search_context_size_high": 0.014 }, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06 }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -27115,7 +27259,7 @@ "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "output_cost_per_video_token": 1.75e-05, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions" ], @@ -27148,7 +27292,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27211,7 +27355,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27268,7 +27412,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -27325,7 +27469,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -28783,6 +28927,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_system_messages": true, @@ -28791,12 +28936,15 @@ "gpt-3.5-turbo-0125": { "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -28806,12 +28954,15 @@ "gpt-3.5-turbo-1106": { "deprecation_date": "2026-09-28", "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 2e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -28839,7 +28990,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 2e-06 + "output_cost_per_token": 2e-06, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-3.5-turbo-instruct-0914": { "input_cost_per_token": 1.5e-06, @@ -28894,12 +29046,15 @@ "gpt-4-0613": { "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 8192, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-05, + "output_cost_per_token_batches": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_system_messages": true, @@ -28940,12 +29095,15 @@ "gpt-4-turbo-2024-04-09": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3e-05, + "output_cost_per_token_batches": 1.5e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -28989,6 +29147,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29031,6 +29190,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29073,6 +29233,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29115,6 +29276,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29153,6 +29315,7 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29190,6 +29353,7 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_priority": 8e-07, "output_cost_per_token_batches": 2e-07, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -29226,6 +29390,7 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29248,6 +29413,7 @@ "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 2.625e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29270,6 +29436,7 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29293,6 +29460,7 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29367,6 +29535,7 @@ "mode": "chat", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29403,6 +29572,7 @@ "mode": "chat", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions" ], @@ -29437,6 +29607,7 @@ "mode": "chat", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29474,6 +29645,7 @@ "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29511,6 +29683,7 @@ "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -29572,7 +29745,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": false, - "deprecation_date": "2027-01-20" + "deprecation_date": "2027-01-20", + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -29600,7 +29774,8 @@ "search_context_size_high": 0.025, "search_context_size_low": 0.025, "search_context_size_medium": 0.025 - } + }, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29621,6 +29796,7 @@ "search_context_size_low": 0.025, "search_context_size_medium": 0.025 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -29769,15 +29945,18 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "input_cost_per_second": 5e-05, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-tts": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_second": 0.00025, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -29909,7 +30088,9 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "input_cost_per_second": 0.0001, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-image-1.5": { "cache_read_input_token_cost": 1.25e-06, @@ -29919,7 +30100,10 @@ "mode": "image_generation", "output_cost_per_token": 1e-05, "input_cost_per_image_token": 8e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations" ], @@ -29934,7 +30118,10 @@ "mode": "image_generation", "output_cost_per_token": 1e-05, "input_cost_per_image_token": 8e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations" ], @@ -29947,7 +30134,9 @@ "litellm_provider": "openai", "mode": "image_generation", "input_cost_per_image_token": 8e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -30394,6 +30583,7 @@ "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -30402,6 +30592,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "search_context_cost_per_query": { @@ -30409,6 +30600,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -30438,6 +30630,7 @@ }, "gpt-5.1": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, @@ -30478,11 +30671,17 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_flex": 6.25e-07, + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_flex": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": false, "supports_minimal_reasoning_effort": false }, "gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, @@ -30523,6 +30722,11 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_flex": 6.25e-07, + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_flex": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": false, "supports_minimal_reasoning_effort": false }, @@ -30574,6 +30778,7 @@ }, "gpt-5.2": { "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_flex": 8.75e-08, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, @@ -30615,11 +30820,17 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 8.75e-07, + "input_cost_per_token_flex": 8.75e-07, + "output_cost_per_token_batches": 7e-06, + "output_cost_per_token_flex": 7e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, "gpt-5.2-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_flex": 8.75e-08, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, @@ -30661,6 +30872,11 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "input_cost_per_token_batches": 8.75e-07, + "input_cost_per_token_flex": 8.75e-07, + "output_cost_per_token_batches": 7e-06, + "output_cost_per_token_flex": 7e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -30754,17 +30970,20 @@ }, "gpt-5.2-pro": { "input_cost_per_token": 2.1e-05, + "input_cost_per_token_batches": 1.05e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.000168, + "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -30793,17 +31012,20 @@ }, "gpt-5.2-pro-2025-12-11": { "input_cost_per_token": 2.1e-05, + "input_cost_per_token_batches": 1.05e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.000168, + "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -30869,6 +31091,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31005,6 +31228,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31073,6 +31297,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31140,6 +31365,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -31203,7 +31429,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://developers.openai.com/api/docs/models/gpt-5.6-cyber", + "source": "https://developers.openai.com/api/docs/pricing", "supports_computer_use": true, "supports_parallel_function_calling": true }, @@ -31373,7 +31599,7 @@ "reasoning_effort_levels": [ "medium" ], - "source": "https://developers.openai.com/api/docs/models/chat-latest", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -31452,7 +31678,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, @@ -31509,7 +31736,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.5-pro": { "input_cost_per_token": 3e-05, @@ -31532,6 +31760,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -31580,6 +31809,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -31657,7 +31887,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -31709,7 +31940,8 @@ "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, - "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-pro": { "input_cost_per_token": 3e-05, @@ -31758,7 +31990,8 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, - "output_cost_per_token_above_272k_tokens_flex": 0.000135 + "output_cost_per_token_above_272k_tokens_flex": 0.000135, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-pro-2026-03-05": { "input_cost_per_token": 3e-05, @@ -31807,7 +32040,8 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, - "output_cost_per_token_above_272k_tokens_flex": 0.000135 + "output_cost_per_token_above_272k_tokens_flex": 0.000135, + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -31858,6 +32092,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -31910,6 +32145,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -31959,6 +32195,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -32008,6 +32245,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", + "source": "https://developers.openai.com/api/docs/pricing", "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, @@ -32026,6 +32264,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -32068,6 +32307,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -32100,6 +32340,7 @@ "cache_read_input_token_cost_priority": 2.5e-07, "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -32108,6 +32349,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "search_context_cost_per_query": { @@ -32115,6 +32357,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32439,6 +32682,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses" ], @@ -32469,6 +32713,7 @@ "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", @@ -32477,6 +32722,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, "search_context_cost_per_query": { @@ -32484,6 +32730,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32517,6 +32764,7 @@ "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", @@ -32525,6 +32773,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, "search_context_cost_per_query": { @@ -32532,6 +32781,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32563,6 +32813,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, + "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -32571,12 +32822,14 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32609,6 +32862,7 @@ "cache_read_input_token_cost_flex": 2.5e-09, "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, + "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", @@ -32617,12 +32871,14 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -32655,9 +32911,11 @@ "deprecation_date": "2026-10-23", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 4e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -32668,9 +32926,11 @@ "deprecation_date": "2026-12-01", "input_cost_per_image_token": 2.5e-06, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 8e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -32690,6 +32950,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32722,6 +32983,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32755,6 +33017,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32790,6 +33053,7 @@ "output_cost_per_token": 2.4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32825,6 +33089,7 @@ "output_cost_per_token": 2.4e-06, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32847,8 +33112,10 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, + "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, "litellm_provider": "openai", "max_input_tokens": 32000, @@ -32857,6 +33124,7 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -32890,6 +33158,7 @@ "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -37609,12 +37878,15 @@ "cache_read_input_token_cost": 7.5e-06, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 6e-05, + "output_cost_per_token_batches": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, @@ -37629,12 +37901,15 @@ "cache_read_input_token_cost": 7.5e-06, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 6e-05, + "output_cost_per_token_batches": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -37656,6 +37931,7 @@ "mode": "responses", "output_cost_per_token": 0.0006, "output_cost_per_token_batches": 0.0003, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37689,6 +37965,7 @@ "mode": "responses", "output_cost_per_token": 0.0006, "output_cost_per_token_batches": 0.0003, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37716,6 +37993,7 @@ "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -37724,6 +38002,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, "search_context_cost_per_query": { @@ -37731,6 +38010,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/chat/completions", @@ -37760,6 +38040,7 @@ "cache_read_input_token_cost_priority": 8.75e-07, "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -37768,6 +38049,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, "search_context_cost_per_query": { @@ -37775,6 +38057,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/chat/completions", @@ -37884,12 +38167,15 @@ "cache_read_input_token_cost": 5.5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_prompt_caching": true, @@ -37902,12 +38188,15 @@ "cache_read_input_token_cost": 5.5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_prompt_caching": true, @@ -37931,6 +38220,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37968,6 +38258,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -37991,10 +38282,11 @@ }, "o4-mini": { "cache_read_input_token_cost": 2.75e-07, - "cache_read_input_token_cost_flex": 1.375e-07, + "cache_read_input_token_cost_flex": 1.38e-07, "cache_read_input_token_cost_priority": 5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, "litellm_provider": "openai", @@ -38003,6 +38295,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "output_cost_per_token_flex": 2.2e-06, "output_cost_per_token_priority": 8e-06, "search_context_cost_per_query": { @@ -38010,6 +38303,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, @@ -38022,10 +38316,11 @@ }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, - "cache_read_input_token_cost_flex": 1.375e-07, + "cache_read_input_token_cost_flex": 1.38e-07, "cache_read_input_token_cost_priority": 5e-07, "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, "litellm_provider": "openai", @@ -38034,6 +38329,7 @@ "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "output_cost_per_token_flex": 2.2e-06, "output_cost_per_token_priority": 8e-06, "search_context_cost_per_query": { @@ -38041,6 +38337,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, @@ -43203,7 +43500,8 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_cost_per_token_batches": 0.0, - "output_vector_size": 3072 + "output_vector_size": 3072, + "source": "https://developers.openai.com/api/docs/pricing" }, "text-embedding-3-small": { "input_cost_per_token": 2e-08, @@ -43214,7 +43512,8 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_cost_per_token_batches": 0.0, - "output_vector_size": 1536 + "output_vector_size": 1536, + "source": "https://developers.openai.com/api/docs/pricing" }, "text-embedding-ada-002": { "input_cost_per_token": 1e-07, @@ -43223,7 +43522,8 @@ "max_tokens": 8191, "mode": "embedding", "output_cost_per_token": 0.0, - "output_vector_size": 1536 + "output_vector_size": 1536, + "source": "https://developers.openai.com/api/docs/pricing" }, "text-embedding-ada-002-v2": { "input_cost_per_token": 1e-07, @@ -43394,7 +43694,7 @@ "input_cost_per_token": 1.2e-06, "output_cost_per_token": 1.2e-06, "max_input_tokens": 131072, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo": { "litellm_provider": "together_ai", @@ -43406,7 +43706,7 @@ "input_cost_per_token": 3e-07, "output_cost_per_token": 3e-07, "max_input_tokens": 32768, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { "deprecation_date": "2026-07-10", @@ -43453,7 +43753,7 @@ "max_input_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://www.together.ai/models/qwen3-coder-480b-a35b-instruct", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43515,7 +43815,7 @@ }, "mode": "chat", "output_cost_per_token": 1.7e-06, - "source": "https://www.together.ai/models/deepseek-v3-1", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43539,7 +43839,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.04e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43573,6 +43873,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 5.9e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43595,6 +43896,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 8.8e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43606,6 +43908,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 1.8e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43622,7 +43925,7 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, "max_input_tokens": 32768, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { "deprecation_date": "2026-04-02", @@ -43634,7 +43937,7 @@ "input_cost_per_token": 1e-07, "output_cost_per_token": 3e-07, "max_input_tokens": 32768, - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { "deprecation_date": "2026-04-16", @@ -43642,6 +43945,7 @@ "litellm_provider": "together_ai", "mode": "chat", "output_cost_per_token": 6e-07, + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43668,7 +43972,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://www.together.ai/models/gpt-oss-120b", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43682,7 +43986,7 @@ "max_input_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-07, - "source": "https://www.together.ai/models/gpt-oss-20b", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43702,7 +44006,7 @@ "max_input_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.1e-06, - "source": "https://www.together.ai/models/glm-4-5-air", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43718,7 +44022,7 @@ }, "mode": "chat", "output_cost_per_token": 2.2e-06, - "source": "https://www.together.ai/models/glm-4-6", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43735,7 +44039,7 @@ }, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://www.together.ai/models/glm-4-7", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43783,7 +44087,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43799,7 +44103,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -43813,7 +44117,7 @@ "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 3.6e-06, - "source": "https://www.together.ai/models/qwen3-5-397b-a17b", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43828,7 +44132,7 @@ "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43853,7 +44157,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 2.5e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -43868,7 +44172,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 3e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { @@ -43879,7 +44183,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 7.5e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/Qwen/Qwen3.7-Plus": { @@ -43889,7 +44193,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.28e-06, - "source": "https://docs.together.ai/docs/serverless-models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { "cache_read_input_token_cost": 2.5e-07, @@ -43899,7 +44203,7 @@ "max_tokens": 1010000, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/arize-ai/qwen-2-1.5b-instruct": { @@ -43909,7 +44213,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1e-07, - "source": "https://docs.together.ai/docs/serverless-models" + "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 3e-08, @@ -43919,7 +44223,7 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 2.8e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43951,7 +44255,7 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 3.96e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43976,7 +44280,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 9.7e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -44012,7 +44316,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/moonshotai/Kimi-K2.7-Code": { @@ -44024,7 +44328,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44045,7 +44349,7 @@ "high", "max" ], - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44063,7 +44367,7 @@ "max_tokens": 512288, "mode": "chat", "output_cost_per_token": 3.6e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44089,7 +44393,7 @@ "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 4.05e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44105,7 +44409,7 @@ "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true }, "together_ai/zai-org/GLM-5.2": { @@ -44117,7 +44421,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44134,7 +44438,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44151,7 +44455,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.ai/v1/models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -44164,6 +44468,7 @@ "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", "mode": "audio_speech", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ] @@ -44172,6 +44477,7 @@ "input_cost_per_character": 3e-05, "litellm_provider": "openai", "mode": "audio_speech", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ] @@ -49497,7 +49803,8 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ], - "deprecation_date": "2027-02-26" + "deprecation_date": "2027-02-26", + "source": "https://developers.openai.com/api/docs/pricing" }, "xai/grok-3": { "cache_read_input_token_cost": 2e-07, @@ -52594,10 +52901,11 @@ "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "input_cost_per_token": 0.0, + "input_cost_per_token": 2e-07, "output_cost_per_token": 0.0, "litellm_provider": "fireworks_ai", - "mode": "rerank" + "mode": "rerank", + "source": "https://api.fireworks.ai/v1/serverless/models" }, "fireworks_ai/accounts/fireworks/models/qwen3-vl-235b-a22b-instruct": { "max_tokens": 262144, @@ -52662,7 +52970,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1.6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -54561,12 +54869,13 @@ }, "gpt-4o-mini-tts-2025-03-20": { "deprecation_date": "2026-07-23", - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_second": 0.00025, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -54579,12 +54888,13 @@ ] }, "gpt-4o-mini-tts-2025-12-15": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_second": 0.00025, "output_cost_per_token": 1e-05, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -54599,24 +54909,28 @@ "gpt-4o-mini-transcribe-2025-03-20": { "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", "output_cost_per_token": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions" ] }, "gpt-4o-mini-transcribe-2025-12-15": { "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", "output_cost_per_token": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions" ] @@ -54635,6 +54949,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -54662,6 +54977,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "source": "https://developers.openai.com/api/docs/pricing", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -54689,6 +55005,7 @@ "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -54740,13 +55057,14 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "deprecation_date": "2027-01-20" + "deprecation_date": "2027-01-20", + "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-realtime-whisper": { - "input_cost_per_second": 0.0002833333333333333, + "input_cost_per_second": 0.000283333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://developers.openai.com/api/docs/models/gpt-realtime-whisper", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -54764,7 +55082,7 @@ "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, - "source": "https://platform.openai.com/docs/api-reference/videos", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "text", "image" @@ -54778,7 +55096,7 @@ "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, - "source": "https://platform.openai.com/docs/api-reference/videos", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "text", "image" @@ -54803,11 +55121,15 @@ "chatgpt-image-latest": { "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2026-12-01", - "input_cost_per_image_token": 1e-05, + "input_cost_per_image_token": 8e-06, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "openai", "mode": "image_generation", - "output_cost_per_image_token": 4e-05, + "output_cost_per_image_token": 3.2e-05, + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -57282,7 +57604,7 @@ "input_cost_per_second": 7.5e-05, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://developers.openai.com/api/docs/models/gpt-transcribe", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/audio/transcriptions", "/v1/realtime/transcription_sessions" @@ -57297,10 +57619,10 @@ "supports_audio_input": true }, "gpt-live-transcribe": { - "input_cost_per_second": 0.0002833333333333333, + "input_cost_per_second": 0.000283333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://developers.openai.com/api/docs/models/gpt-live-transcribe", + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -57315,10 +57637,10 @@ "supports_audio_input": true }, "gpt-live-1": { - "input_cost_per_second": 0.0008333333333333334, + "input_cost_per_second": 0.000833333333333, "litellm_provider": "openai", "mode": "realtime", - "source": "https://developers.openai.com/api/docs/models/gpt-live-1", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "text", "audio" @@ -57332,13 +57654,13 @@ "supports_function_calling": true }, "gpt-realtime-translate": { - "input_cost_per_second": 0.0005666666666666667, + "input_cost_per_second": 0.000566666666667, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "max_tokens": 2000, "mode": "realtime", - "source": "https://developers.openai.com/api/docs/models/gpt-realtime-translate", + "source": "https://developers.openai.com/api/docs/pricing", "supported_modalities": [ "audio" ], @@ -57366,7 +57688,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "source": "https://platform.claude.com/docs/en/about-claude/models/overview", + "source": "https://platform.claude.com/docs/en/about-claude/pricing", "supports_adaptive_thinking": true, "thinking_always_on": true, "supports_mid_conversation_system": true, @@ -57427,7 +57749,7 @@ "supports_output_config": true, "prompt_cache_min_tokens": 512, "supports_native_structured_output": true, - "source": "https://platform.claude.com/docs/en/models/mythos-5-1/overview" + "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-mythos-preview": { "cache_creation_input_token_cost": 1.25e-05, @@ -57779,6 +58101,7 @@ }, "vertex_ai/gemini-3.5-live-translate-preview": { "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_second": 8.83333333333e-05, "input_cost_per_token": 3.5e-06, "litellm_provider": "vertex_ai", "mode": "realtime", @@ -57881,14 +58204,17 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -57897,14 +58223,17 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -57920,26 +58249,29 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/kimi-k3": { "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_priority": 3.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_priority": 1.875e-05, "reasoning_effort_levels": [ "low", "high", "max" ], - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -57948,14 +58280,17 @@ }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -57964,14 +58299,17 @@ }, "fireworks_ai/deepseek-v4p1-flash": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "output_cost_per_token_priority": 8.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -57987,7 +58325,7 @@ "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true @@ -58026,19 +58364,22 @@ }, "fireworks_ai/kimi-k3": { "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_priority": 3.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_priority": 1.875e-05, "reasoning_effort_levels": [ "low", "high", "max" ], - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58089,12 +58430,15 @@ }, "fireworks_ai/qwen3p8-max": { "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 9e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58110,7 +58454,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58142,7 +58486,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58158,7 +58502,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58190,7 +58534,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58199,12 +58543,15 @@ }, "fireworks_ai/accounts/fireworks/models/qwen3p8-max": { "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 9e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58220,7 +58567,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6.6e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -58257,7 +58604,7 @@ "high", "max" ], - "source": "https://docs.fireworks.ai/serverless/pricing", + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -60955,14 +61302,17 @@ }, "fireworks_ai/accounts/fireworks/models/glm-5p3": { "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_priority": 3.25e-07, "input_cost_per_token": 1.4e-06, + "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 5.5e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -60971,13 +61321,16 @@ }, "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_priority": 3.75e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 1.875e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.fireworks.ai/serverless/pricing", + "output_cost_per_token_priority": 6.25e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -61005,7 +61358,7 @@ "max_output_tokens": 40960, "max_tokens": 40960, "mode": "embedding", - "source": "https://docs.fireworks.ai/serverless/pricing" + "source": "https://api.fireworks.ai/v1/serverless/models" }, "zai/glm-5.2": { "cache_creation_input_token_cost": 0, @@ -61029,7 +61382,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 4.7e-07, - "source": "https://docs.together.ai/docs/serverless-models" + "source": "https://api.together.ai/v1/models" }, "together_ai/moonshotai/Kimi-K2.6": { "deprecation_date": "2026-08-19", @@ -61039,7 +61392,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/moonshotai/Kimi-K2.5-fp4": { "input_cost_per_token": 5e-07, @@ -61047,7 +61400,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/MiniMaxAI/MiniMax-M2.7": { "input_cost_per_token": 3e-07, @@ -61056,7 +61409,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 196608, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/zai-org/GLM-5": { "deprecation_date": "2026-06-22", @@ -61065,7 +61418,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 202752, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/zai-org/GLM-5.1": { "deprecation_date": "2026-07-10", @@ -61075,7 +61428,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 202752, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-0528": { "input_cost_per_token": 3e-06, @@ -61083,7 +61436,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 163840, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-Coder-Next-FP8": { "deprecation_date": "2026-05-14", @@ -61092,7 +61445,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-VL-32B-Instruct": { "deprecation_date": "2026-02-25", @@ -61101,7 +61454,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen3-VL-8B-Instruct": { "deprecation_date": "2026-04-16", @@ -61110,7 +61463,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Ministral-3-14B-Instruct-2512": { "input_cost_per_token": 2e-07, @@ -61118,7 +61471,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 262144, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { "input_cost_per_token": 6e-08, @@ -61126,7 +61479,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 131072, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/mistralai/Mistral-7B-Instruct-v0.3": { "input_cost_per_token": 2e-07, @@ -61134,7 +61487,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 32768, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/QwQ-32B": { "deprecation_date": "2025-11-13", @@ -61143,7 +61496,7 @@ "litellm_provider": "together_ai", "max_input_tokens": 131072, "mode": "chat", - "source": "https://api.together.xyz/v1/models" + "source": "https://api.together.ai/v1/models" }, "cerebras/gemma-4-31b": { "input_cost_per_token": 9.9e-07, @@ -65270,5 +65623,260 @@ "supports_tool_choice": false, "supports_response_schema": true, "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://api.fireworks.ai/v1/serverless/models" + }, + "together_ai/arcee-ai/trinity-mini": { + "input_cost_per_token": 4.5e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/deepseek-coder-33b-instruct": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "input_cost_per_token": 2e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "input_cost_per_token": 1.6e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models" + }, + "vertex_ai/gemini-2.5-flash-native-audio": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-2.5-flash-preview-tts": { + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_audio_token": 1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.1-flash-live-preview": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_second": 8.33333333333e-05, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.1-flash-tts-preview": { + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.5-transcribe": { + "input_cost_per_audio_token": 2e-06, + "input_cost_per_second": 5e-05, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-3.5-transcribe-live": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_second": 8.33333333333e-05, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-omni-1.1-flash": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 9e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemini-robotics-er-2": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "vertex_ai/gemma-4-26b-a4b-it": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai", + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, + "together_ai/google/gemma-2-27b-it": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "gpt-5.5-cyber": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 1.25e-05, + "litellm_provider": "openai", + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-rosalind-research": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "together_ai/meta-llama/Llama-3-8b-chat-hf": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Llama-3.1-405B-Instruct": { + "input_cost_per_token": 3.5e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 3.5e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Llama-3.2-1B-Instruct": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 6e-08, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Llama-3.2-3B-Instruct": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 6e-08, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Meta-Llama-3-70B-Instruct-Turbo": { + "input_cost_per_token": 8.8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8.8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/meta-llama/Meta-Llama-3-8B-Instruct": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF": { + "input_cost_per_token": 8.8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8.8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2-1.5B-Instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 2e-08, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2-72B-Instruct": { + "input_cost_per_token": 9e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2-VL-72B-Instruct": { + "input_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-14B-Instruct": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-72B-Instruct": { + "input_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-Coder-32B-Instruct": { + "input_cost_per_token": 8e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://api.together.ai/v1/models" + }, + "together_ai/Qwen/Qwen2.5-VL-72B-Instruct": { + "input_cost_per_token": 1.95e-06, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://api.together.ai/v1/models" } } From 4c022a3089cbfbc377cc700db7ade6309f3b1e1f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 23:34:12 -0700 Subject: [PATCH 062/297] feat(pricing): add azure gpt-chat-latest global and data zone rates --- ...odel_prices_and_context_window_backup.json | 111 ++++++++++++++++++ model_prices_and_context_window.json | 111 ++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 31 +++++ .../test_litellm/test_model_prices_schema.py | 8 ++ 4 files changed, 261 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2220d0e1fe5..86605d77aab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7409,6 +7409,80 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure/chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/us/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, @@ -7675,6 +7749,43 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-chat-latest": { + "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/eu/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2220d0e1fe5..86605d77aab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7409,6 +7409,80 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure/chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/us/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, @@ -7675,6 +7749,43 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-chat-latest": { + "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/eu/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index cbe6fe198c9..3c6a13a7f01 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2101,6 +2101,37 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( assert completion_cost == pytest.approx(zone_multiplier * output_multiplier * completion_tokens * 5e-5) +@pytest.mark.parametrize( + "model,input_rate,cache_read_rate,output_rate", + [ + ("azure/gpt-chat-latest", 5e-6, 5e-7, 3e-5), + ("azure/chat-latest", 5e-6, 5e-7, 3e-5), + ("azure/us/gpt-chat-latest", 5.5e-6, 5.5e-7, 3.3e-5), + ], +) +def test_generic_cost_per_token_azure_gpt_chat_latest_price_sheet( + _local_model_cost_map, model, input_rate, cache_read_rate, output_rate +): + """The Azure OpenAI price sheet lists GPT-Chat Latest at $5 input, $0.50 cached input and $30 output per 1M + tokens on Global, and $5.50, $0.55 and $33 on Data Zone. Foundry names the product gpt-chat-latest and the + OpenAI API names the same model chat-latest, so both spellings bill the Global sheet. + """ + prompt_tokens = 100000 + cached_tokens = 40000 + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="azure") + + assert prompt_cost == pytest.approx((prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate) + assert completion_cost == pytest.approx(completion_tokens * output_rate) + + def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 0b9dbd23097..e562797fbe8 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -221,6 +221,14 @@ def test_chat_latest_declares_the_one_effort_openai_accepts(prices: dict): assert resolve_supported_reasoning_efforts(prices["chat-latest"], deployment_is_mapped=True) == ("medium",) +@pytest.mark.parametrize("key", ["azure/gpt-chat-latest", "azure/chat-latest", "azure/us/gpt-chat-latest"]) +def test_azure_gpt_chat_latest_declares_the_one_effort_azure_accepts(prices: dict, key: str): + """Azure answers every reasoning_effort on a gpt-chat-latest deployment except medium with + "Unsupported value ... Supported values are: 'medium'", the same fixed level OpenAI's chat-latest + carries, so the Foundry product name and the OpenAI API name both declare that one level.""" + assert resolve_supported_reasoning_efforts(prices[key], deployment_is_mapped=True) == ("medium",) + + BEDROCK_OPENAI_GPT_MARKERS: Final = ("openai.gpt-5.4", "openai.gpt-5.5", "openai.gpt-5.6", "openai.gpt-6-astra") BEDROCK_PROVIDERS: Final = frozenset(("bedrock", "bedrock_converse", "bedrock_mantle")) BEDROCK_ROW_PREFIXES: Final = ("bedrock_mantle/", "us.", "global.") From 09314f239c34f35a86ce4155b8c4d445944a6038 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 00:54:29 -0700 Subject: [PATCH 063/297] fix(guardrails): hand per-message rewrites back as structured_messages A guardrail that rewrites text per chat message now returns the rewritten rows as structured_messages instead of only texts, so the Responses and chat handlers write the rewrite back through the structured path. The generic guardrail API response accepts an optional structured_messages list, Prompt Security modify builds one from modified_messages, and rows a server echoes back exactly as shown are restored to the original row objects because the request model drops undeclared keys. Texts-only per-message answers keep the named rejection on both endpoints. --- .../base_llm/guardrail_translation/utils.py | 42 ++--- .../guardrail_translation/handler.py | 27 +--- .../generic_guardrail_api.py | 30 +++- .../prompt_security/prompt_security.py | 8 +- .../guardrail_hooks/generic_guardrail_api.py | 15 +- ...test_openai_responses_guardrail_handler.py | 148 +++++++++--------- .../test_generic_guardrail_api.py | 105 +++++++++++++ .../test_prompt_security_guardrails.py | 14 +- 8 files changed, 244 insertions(+), 145 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 383e668e45c..1172c93959b 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles @@ -390,7 +389,7 @@ def _part_with_text(part: object, text: str) -> object: return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts -def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]: +def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]: text_part_indices: Final = tuple( index for index, part in enumerate(content) if _content_part_text(part) is not None ) @@ -401,39 +400,18 @@ def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> ] -def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues: +def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: + """Swap one rewritten text into each text slot of a chat row, in order. + + A slot is a string ``content`` or one list part carrying a string ``text``; + images and other parts ride along untouched. Returns None unless the counts + line up exactly, so a rewrite never lands on the wrong slot. + """ + if message_text_slot_count(message) != len(texts): + return None content: Final = message.get("content") if not isinstance(content, (str, list)): return message rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped - - -def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: - if message_text_slot_count(message) != len(texts): - return None - return _message_with_slot_texts(message, texts) - - -def messages_with_slot_texts( - messages: Sequence[AllMessageValues], - texts: Sequence[str], -) -> list[AllMessageValues] | None: - """Spread one flat list of rewritten texts over the messages' text slots, in order. - - A slot is a string ``content`` or one list part carrying a string ``text``; - images and other parts ride along untouched. A guardrail that answers one - text per message it saw produces exactly this shape, which stops matching - the endpoint's own per-text extraction as soon as the request carries - instructions or tool items. Returns None unless the counts line up exactly, - so a rewrite never lands on the wrong slot. - """ - slot_counts: Final = tuple(message_text_slot_count(message) for message in messages) - if sum(slot_counts) != len(texts): - return None - offsets: Final = tuple(accumulate(slot_counts, initial=0)) - return [ # mutable-ok: guardrail rows travel as a list - _message_with_slot_texts(message, texts[start:end]) - for message, start, end in zip(messages, offsets, offsets[1:]) - ] diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 61903a54cd3..2fe11d9f7bd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -53,7 +53,6 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, - messages_with_slot_texts, stream_item_field, stream_item_fingerprint, stream_item_items, @@ -396,20 +395,6 @@ def _patched_request_fields( ) -def _guardrailed_structured_messages( - structured_messages: Sequence[AllMessageValues] | None, - sent_text_count: int, - guardrailed_inputs: GenericGuardrailAPIInputs, -) -> Sequence[AllMessageValues] | None: - returned: Final = guardrailed_inputs.get("structured_messages") - if returned is not None and returned is not structured_messages: - return returned - rewritten_texts: Final = guardrailed_inputs.get("texts") - if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count: - return None - return messages_with_slot_texts(structured_messages, rewritten_texts) - - def _patch_or_convert_request_fields( raw_input: object, instructions: object, @@ -488,8 +473,7 @@ class OpenAIResponsesHandler(BaseTranslation): form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) - sent_texts: Final = extracted.inputs.get("texts") - if not sent_texts: + if not extracted.inputs.get("texts"): return data if structured_messages: extracted.inputs["structured_messages"] = structured_messages @@ -502,9 +486,7 @@ class OpenAIResponsesHandler(BaseTranslation): self._apply_guardrailed_tools_to_data( data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) - written_back: Final = self._written_back_request_fields( - data, structured_messages, len(sent_texts), guardrailed_inputs - ) + written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) if written_back is not None: data["input"] = list(written_back.input) # mutable-ok: JSON body if written_back.instructions is None: @@ -571,11 +553,10 @@ class OpenAIResponsesHandler(BaseTranslation): def _written_back_request_fields( data: Mapping[str, object], structured_messages: Sequence[AllMessageValues] | None, - sent_text_count: int, guardrailed_inputs: GenericGuardrailAPIInputs, ) -> _RequestFields | None: - guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs) - if guardrailed is None: + guardrailed: Final = guardrailed_inputs.get("structured_messages") + if guardrailed is None or guardrailed is structured_messages: return None return _patch_or_convert_request_fields( data.get("input"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index d8296003ae9..16159d32a7f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -7,7 +7,7 @@ import fnmatch import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional import httpx @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, @@ -150,6 +150,22 @@ def _extract_inbound_headers( return None +def _rows_with_unchanged_originals( + original_rows: Sequence[AllMessageValues] | None, + shown_rows: Sequence[AllMessageValues] | None, + returned_rows: Sequence[AllMessageValues], +) -> tuple[AllMessageValues, ...]: + """The request model drops row keys its message types do not declare, so a + row the server echoes back verbatim is restored to the original row object; + only rows the server actually changed reach the endpoint write-back.""" + if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): + return tuple(returned_rows) + return tuple( + original if returned == shown else returned + for original, shown, returned in zip(original_rows, shown_rows, returned_rows) + ) + + class GenericGuardrailAPI(CustomGuardrail): """ Generic Guardrail API integration for LiteLLM. @@ -322,6 +338,8 @@ class GenericGuardrailAPI(CustomGuardrail): texts: list, images: list[str] | None, tools: list[ChatCompletionToolParam] | None, + structured_messages: Sequence[AllMessageValues] | None, + shown_messages: Sequence[AllMessageValues] | None, guardrail_response: GenericGuardrailAPIResponse, ) -> GenericGuardrailAPIInputs: # Action is NONE or no modifications needed @@ -336,6 +354,12 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + if guardrail_response.structured_messages: + return_inputs["structured_messages"] = list( # mutable-ok: guardrail inputs take a list + _rows_with_unchanged_originals( + structured_messages, shown_messages, guardrail_response.structured_messages + ) + ) if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs @@ -473,6 +497,8 @@ class GenericGuardrailAPI(CustomGuardrail): texts=texts, images=images, tools=tools, + structured_messages=structured_messages, + shown_messages=guardrail_request.structured_messages, guardrail_response=guardrail_response, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 72c87a793a5..41d3b202344 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -287,7 +287,7 @@ class PromptSecurityGuardrail(CustomGuardrail): structured_messages, modified_messages ) if rewritten_messages is not None: - inputs["structured_messages"] = rewritten_messages + inputs["structured_messages"] = list(rewritten_messages) # mutable-ok: guardrail inputs take a list return inputs @@ -298,7 +298,7 @@ class PromptSecurityGuardrail(CustomGuardrail): self, structured_messages: Sequence[AllMessageValues], modified_messages: Sequence[Mapping[str, object]], - ) -> list[AllMessageValues] | None: + ) -> tuple[AllMessageValues, ...] | None: sent_indices: Final = tuple( index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message) ) @@ -313,9 +313,7 @@ class PromptSecurityGuardrail(CustomGuardrail): ) if len(replacements) != len(sent_indices): return None - return [ # mutable-ok: guardrail inputs take a list - replacements.get(index, message) for index, message in enumerate(structured_messages) - ] + return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages)) async def _apply_guardrail_on_response( self, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4a868c48352..44e2cc2404f 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,4 +1,5 @@ -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict @@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int: return 0 +def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None: + if not isinstance(value, list): + return None + if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value): + return None + return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" texts: list[str] | None images: list[str] | None tools: list[GuardrailToolParam] | None + structured_messages: Sequence[AllMessageValues] | None action: str blocked_reason: str | None stream_holdback_chars: list[int] | None @@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse: images: list[str] | None = None, tools: list[GuardrailToolParam] | None = None, stream_holdback_chars: list[int] | None = None, + structured_messages: Sequence[AllMessageValues] | None = None, ) -> None: self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images self.tools = tools + self.structured_messages = structured_messages # Number of trailing chars, indexed the same as ``texts``, that the # framework must withhold from streaming emission until the next # processing round (word-boundary safety for text transformations). @@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse: images=data.get("images"), tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, + structured_messages=structured_messages_from_response(data.get("structured_messages")), ) 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 b4b467773da..394f134e99c 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 @@ -8,7 +8,7 @@ with guardrail transformations. import copy from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import logging @@ -31,6 +31,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -2342,103 +2343,94 @@ SSN = "123-45-6789" REDACTED_SSN = "" -def _slot_texts(message: dict) -> list[str]: - content = message.get("content") - if isinstance(content, str): - return [content] - if isinstance(content, list): - return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)] - return [] +def _redacted(value: object) -> object: + if isinstance(value, str): + return value.replace(SSN, REDACTED_SSN) + if isinstance(value, list): + return [{**part, "text": _redacted(part["text"])} if "text" in part else part for part in value] + return value -class PerMessageRedactionGuardrail(CustomGuardrail): - """Guardrail that answers one redacted text per message it was shown and hands - back only texts, the way Prompt Security in modify mode and a generic guardrail - API server that scans per message do.""" +def _per_message_guardrail_server(structured_messages_in_answer: bool) -> Callable[..., MagicMock]: + """Answers one redacted text per chat row it was shown, the way a guardrail + that scans per message does, and optionally the rewritten rows themselves.""" - def __init__(self, extra_texts: int = 0): - super().__init__(guardrail_name="per-message-redactor") - self.extra_texts = extra_texts + def post(url: str, json: dict, headers: dict) -> MagicMock: + rows = json["structured_messages"] + answer: dict = { + "action": "GUARDRAIL_INTERVENED", + "texts": [_redacted(row["content"]) if isinstance(row.get("content"), str) else "" for row in rows], + } + if structured_messages_in_answer: + answer["structured_messages"] = [{**row, "content": _redacted(row.get("content"))} for row in rows] + response = MagicMock() + response.json.return_value = answer + response.raise_for_status = MagicMock() + return response - async def apply_guardrail( - self, - inputs: GenericGuardrailAPIInputs, - request_data: dict, - input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, - ) -> GenericGuardrailAPIInputs: - messages = inputs.get("structured_messages") or [] - texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)] - return {**inputs, "texts": texts + ["junk"] * self.extra_texts} + return post -class TestPerMessageTextWriteBack: - """A guardrail that rewrites one text per message it saw must land on the - instructions and the input items those messages came from, not be rejected.""" +def _per_message_redactor() -> GenericGuardrailAPI: + return GenericGuardrailAPI( + api_base="https://guardrail.test", + guardrail_name="per-message-redactor", + event_hook="pre_call", + default_on=True, + ) + + +def _tool_replay_request() -> dict: + return { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + {"type": "function_call", "call_id": "call_1", "name": "lookup_customer", "arguments": '{"id": "42"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ], + } + + +class TestPerMessageRewriteWriteBack: + """A guardrail that rewrites per chat row hands the rows back as + structured_messages, and the handler lands them on the instructions and the + input items they came from; the same rewrite handed back as texts alone has + no item to land on and is rejected by name instead of sent unrewritten.""" @pytest.mark.asyncio - async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self): - handler = OpenAIResponsesHandler() - function_call_item = { - "type": "function_call", - "call_id": "call_1", - "name": "lookup_customer", - "arguments": '{"query": "' + SSN + '"}', - } - data = { - "model": "gpt-5.6", - "instructions": "Never repeat the SSN " + SSN + " back.", - "input": [ - {"role": "user", "content": "Look up " + SSN + " for me."}, - function_call_item, - {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, - ], - } + async def test_structured_rows_land_on_instructions_and_tool_output(self): + guardrail = _per_message_redactor() + data = _tool_replay_request() + function_call_item = data["input"][1] - result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." - assert [item.get("type", item.get("role")) for item in result["input"]] == [ - "user", - "function_call", - "function_call_output", - ] - assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] + assert _texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] assert result["input"][1] == function_call_item - assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}' - assert result["input"][2]["call_id"] == "call_1" - - @pytest.mark.asyncio - async def test_string_input_with_instructions_keeps_the_two_apart(self): - handler = OpenAIResponsesHandler() - data = { - "model": "gpt-5.6", - "instructions": "Redact " + SSN + " everywhere.", - "input": "My SSN is " + SSN + ".", + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_1", + "output": '{"ssn": "' + REDACTED_SSN + '"}', } - result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) - - assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere." - assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] - @pytest.mark.asyncio - async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self): + async def test_texts_only_per_message_answer_is_rejected_by_name(self): from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - handler = OpenAIResponsesHandler() - original_input = [ - {"role": "user", "content": "Look up " + SSN + " for me."}, - {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, - ] - data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)} + guardrail = _per_message_redactor() + data = _tool_replay_request() + original = copy.deepcopy(data) - with pytest.raises(UnappliableRequestRewrite) as excinfo: - await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1)) + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)): + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await OpenAIResponsesHandler().process_input_messages(data, guardrail) assert excinfo.value.guardrail_name == "per-message-redactor" - assert data["input"] == original_input - assert data["instructions"] == "Be terse." + assert data["input"] == original["input"] + assert data["instructions"] == original["instructions"] class TestProvenancePatching: 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 83cc9ae8bb9..cc9942e0e40 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 @@ -582,6 +582,111 @@ class TestGuardrailActions: assert result_images is None +class TestStructuredMessagesInResponse: + """A guardrail server that rewrites per chat row answers with the rewritten + rows as structured_messages, which the endpoint handlers write back by row.""" + + @pytest.mark.asyncio + async def test_returned_rows_are_handed_back_as_structured_messages( + self, generic_guardrail, mock_request_data_input + ): + rewritten_rows = [ + {"role": "system", "content": "Never repeat an SSN."}, + {"role": "user", "content": "Look up [REDACTED] for me."}, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'}, + ] + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Never repeat an SSN.", "Look up [REDACTED] for me.", '{"ssn": "[REDACTED]"}'], + "structured_messages": rewritten_rows, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert guardrailed_inputs["structured_messages"] == rewritten_rows + assert guardrailed_inputs["texts"] == mock_response.json.return_value["texts"] + + @pytest.mark.asyncio + async def test_rows_echoed_back_as_shown_keep_their_original_keys( + self, generic_guardrail, mock_request_data_input + ): + tool_call_row = { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0} + ], + } + original_rows = [ + {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"}, + tool_call_row, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + ] + + def echo_with_tool_output_redacted(url, json, headers): + shown_rows = json["structured_messages"] + assert "index" not in shown_rows[1]["tool_calls"][0] + assert "name" not in shown_rows[0] + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Look up 123-45-6789 for me."], + "structured_messages": [ + shown_rows[0], + shown_rows[1], + {**shown_rows[2], "content": '{"ssn": "[REDACTED]"}'}, + ], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_with_tool_output_redacted): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."], "structured_messages": original_rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + returned_rows = guardrailed_inputs["structured_messages"] + assert returned_rows[0] is original_rows[0] + assert returned_rows[1] is tool_call_row + assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "structured_messages", + [[], [{"content": "a row with no role"}], "not a list"], + ids=["empty", "no_role", "not_a_list"], + ) + async def test_rows_that_are_not_chat_messages_are_ignored( + self, generic_guardrail, mock_request_data_input, structured_messages + ): + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["[REDACTED]"], + "structured_messages": structured_messages, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["[REDACTED]"] + + class TestImageSupport: """Test image handling in guardrail requests""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 9e83098eb04..70083e50f01 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -1,5 +1,6 @@ import asyncio import base64 +from collections.abc import Mapping, Sequence from unittest.mock import AsyncMock, patch import pytest @@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im PromptSecurityGuardrailMissingSecrets, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.llms.openai import AllMessageValues def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): @@ -174,7 +176,7 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] -def _modify_response(modified_messages: list) -> Response: +def _modify_response(modified_messages: Sequence[Mapping[str, object]]) -> Response: mock_response = Response( json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}}, status_code=200, @@ -184,7 +186,7 @@ def _modify_response(modified_messages: list) -> Response: return mock_response -def _tool_replay_messages() -> list: +def _tool_replay_messages() -> list[AllMessageValues]: return [ {"role": "system", "content": "Never echo an SSN like 123-45-6789."}, { @@ -224,7 +226,9 @@ async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatc ] with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) assert result["structured_messages"] == [ {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, @@ -257,7 +261,9 @@ async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}] with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) assert result["structured_messages"] is messages assert result["texts"] == ["Look up [REDACTED]"] From 76ae35dfcdefbb6b09cc576909c69b54bae5999d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:14:35 -0700 Subject: [PATCH 064/297] fix(types): break the CachedTokensDetails import cycle CodeQL flagged two module-level cyclic imports introduced by defining CachedTokensDetails in litellm.types.llms.openai and importing it from litellm.types.utils and litellm.cost_calculator. The class now lives in litellm.types.llms.base, which imports nothing from litellm, and every user imports it from there. Also pins that combining realtime usages where only one response.done carries cached_tokens_details keeps the earlier modality split in both orders, and commits the regenerated dashboard API types. --- litellm/cost_calculator.py | 2 +- .../transformation.py | 2 +- litellm/types/llms/base.py | 6 +++ litellm/types/llms/openai.py | 8 +--- litellm/types/utils.py | 2 +- tests/test_litellm/test_cost_calculator.py | 43 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 7 files changed, 53 insertions(+), 12 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ce5c84907a1..f5319776213 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -97,6 +97,7 @@ from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_ro from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( HttpxBinaryResponseContent, ImageGenerationRequestQuality, @@ -109,7 +110,6 @@ from litellm.types.llms.openai import ( ) from litellm.types.rerank import RerankBilledUnits, RerankResponse from litellm.types.utils import ( - CachedTokensDetails, CallTypesLiteral, LiteLLMRealtimeStreamLoggingObject, LlmProviders, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7a215cdbb12..64324c6cad8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -43,9 +43,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( AllMessageValues, - CachedTokensDetails, ChatCompletionAssistantMessage, ChatCompletionImageObject, ChatCompletionImageUrlObject, diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index f09727ad92b..938aa8064c9 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -75,3 +75,9 @@ class HiddenParams(OpenAIObject): data: Final = super().model_dump(**kwargs) data["_response_ms"] = self._response_ms return data + + +class CachedTokensDetails(BaseModel): + text_tokens: int | None = None + audio_tokens: int | None = None + image_tokens: int | None = None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 7a86c27efae..9bdad700d4b 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -82,7 +82,7 @@ from typing_extensions import ( override, ) -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject +from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, CachedTokensDetails from litellm.types.responses.main import ( CustomToolCallOutputItem, GenericResponseOutputItem, @@ -1285,12 +1285,6 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} -class CachedTokensDetails(BaseModel): - text_tokens: int | None = None - audio_tokens: int | None = None - image_tokens: int | None = None - - class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 07e6fc5838b..1d73542c9bb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -48,6 +48,7 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.types.llms.base import ( BaseLiteLLMOpenAIResponseObject, + CachedTokensDetails, LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo @@ -60,7 +61,6 @@ from .llms.base import HiddenParams from .llms.openai import ( AllMessageValues, Batch, - CachedTokensDetails, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionRedactedThinkingBlock, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a16f1b8fe4a..68e9b6143a0 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -19,6 +19,7 @@ from litellm.cost_calculator import ( ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import OpenAIRealtimeStreamList from litellm.types.rerank import RerankResponse from litellm.types.utils import ( @@ -4896,6 +4897,48 @@ def test_realtime_combine_sums_nested_cached_tokens_details(): assert combined.prompt_tokens_details.cached_tokens_details.image_tokens is None +@pytest.mark.parametrize("details_first", [True, False]) +def test_realtime_combine_keeps_cached_split_when_only_one_usage_has_details(details_first: bool): + with_details: Final = { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 283, + "output_tokens": 0, + "total_tokens": 283, + "input_token_details": { + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + } + }, + } + without_details: Final = { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 150, + "output_tokens": 0, + "total_tokens": 150, + "input_token_details": {"text_tokens": 50, "audio_tokens": 100, "cached_tokens": 100}, + } + }, + } + results: OpenAIRealtimeStreamList = ( + [with_details, without_details] if details_first else [without_details, with_details] + ) + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.cached_tokens == 292 + assert combined.prompt_tokens_details.cached_tokens_details == CachedTokensDetails(text_tokens=64, audio_tokens=128) + + def test_usage_without_cached_tokens_details_omits_key(): usage = Usage( prompt_tokens=10, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 93f3911b32e04c124027e8c5a12961fdcb0fca90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 065/297] fix(guardrails): drop the types import CodeQL reads as a package cycle --- litellm/llms/base_llm/guardrail_translation/utils.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 1172c93959b..34d648cf184 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Mapping, Sequence -from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles from pydantic import BaseModel @@ -390,13 +389,10 @@ def _part_with_text(part: object, text: str) -> object: def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]: - text_part_indices: Final = tuple( - index for index, part in enumerate(content) if _content_part_text(part) is not None - ) - replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts))) + remaining_texts: Final = iter(texts) return [ # mutable-ok: message content stays a JSON list - _part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part - for index, part in enumerate(content) + _part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part + for part in content ] From 1760fe628436fac7fb292fae9ecdf6016a51fc55 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 066/297] refactor(guardrails): return a fresh inputs mapping from the Prompt Security modify branch --- .../prompt_security/prompt_security.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 41d3b202344..5d533e98ff6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -37,6 +37,18 @@ class PromptSecurityGuardrailMissingSecrets(Exception): pass +def _inputs_with_modifications( + inputs: GenericGuardrailAPIInputs, + modified_texts: list[str], + rewritten_messages: Sequence[AllMessageValues] | None, +) -> GenericGuardrailAPIInputs: + texts_patch: Final[GenericGuardrailAPIInputs] = {"texts": modified_texts} if modified_texts else {} + messages_patch: Final[GenericGuardrailAPIInputs] = ( + {"structured_messages": list(rewritten_messages)} if rewritten_messages is not None else {} + ) + return {**inputs, **texts_patch, **messages_patch} + + class _ProtectVerdict(TypedDict, total=False): """One side (``prompt`` or ``response``) of an ``/api/protect`` verdict.""" @@ -280,14 +292,11 @@ class PromptSecurityGuardrail(CustomGuardrail): ) elif action == "modify": modified_messages: Final = result.get("modified_messages", []) - modified_texts: Final = self._extract_texts_from_messages(modified_messages) - if modified_texts: - inputs["texts"] = modified_texts - rewritten_messages: Final = self._structured_messages_with_modifications( - structured_messages, modified_messages + return _inputs_with_modifications( + inputs, + self._extract_texts_from_messages(modified_messages), + self._structured_messages_with_modifications(structured_messages, modified_messages), ) - if rewritten_messages is not None: - inputs["structured_messages"] = list(rewritten_messages) # mutable-ok: guardrail inputs take a list return inputs From 23a98cb85134342e20ccb7387875c9b0037bc7d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 067/297] chore(ui): regenerate schema.d.ts after merging main --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From cccff657cf376843a1e9c732e9ec33c46183e01e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:35:19 -0700 Subject: [PATCH 068/297] fix(guardrails): scan the Anthropic top-level system prompt and tool_use arguments --- .../chat/guardrail_translation/handler.py | 224 +++++++++++++--- .../test_anthropic_guardrail_handler.py | 248 ++++++++++++++++-- 2 files changed, 413 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..eb16278c9bd 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -103,9 +103,24 @@ class ToolResultBlockTextTarget: block_idx: int -InputWriteBackTarget = ( - MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget -) +@dataclass(frozen=True, slots=True) +class SystemStringTarget: + pass + + +@dataclass(frozen=True, slots=True) +class SystemBlockTextTarget: + block_idx: int + + +@dataclass(frozen=True, slots=True) +class ToolUseInputTarget: + msg_idx: int + content_idx: int + + +MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget +InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: @@ -146,10 +161,17 @@ class ScannedText: target: InputWriteBackTarget +@dataclass(frozen=True, slots=True) +class ScannedToolCall: + tool_call: ChatCompletionToolCallChunk + target: ToolUseInputTarget + + @dataclass(frozen=True, slots=True) class ExtractedInput: scanned: tuple[ScannedText, ...] images: tuple[str, ...] + tool_calls: tuple[ScannedToolCall, ...] = () EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) @@ -161,6 +183,76 @@ class _ToolCallShape: arguments: str +def _is_client_tool_use(block: Mapping[str, object]) -> bool: + return ( + block.get("type") == "tool_use" + and isinstance(block.get("id"), str) + and isinstance(block.get("name"), str) + and isinstance(block.get("input"), Mapping) + ) + + +def _write_back_system_block(system: object, block_idx: int, response: str) -> None: + if not isinstance(system, list): + return + text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text") + if block_idx < len(text_blocks): + text_blocks[block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + + +def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None: + content: Final = message.get("content", None) + if content is None: + return + match target: + case MessageContentTarget(): + if isinstance(content, str): + message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place + case ContentBlockTextTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultStringTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["content"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): + if isinstance(content, list): + content[content_idx]["content"][block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case _: + assert_never(target) + + +def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: + content: Final = message.get("content", None) + block: Final = content[target.content_idx] if isinstance(content, list) else None + if not isinstance(block, dict): + return + try: + rewritten_input: Final = json.loads(shape.arguments) + except json.JSONDecodeError: + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + if not isinstance(rewritten_input, dict): + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place + if shape.name is not None and shape.name != block.get("name"): + block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place + + @dataclass(frozen=True, slots=True) class _SSEFieldRewrite: """One field of one nested section of a buffered SSE event, rewritten.""" @@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation): skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply) scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply) - # Exclude only the trusted top-level prompt. In-sequence system entries are untrusted - # and must stay aligned with texts_to_check for positional masking. When the top-level - # prompt is included, the pre-existing count mismatch disables positional masking. + # The top-level prompt is translated on its own below so it can be hoisted in front of + # any mid-turn system entries and scanned first, aligned with that structured position. translation_source: Final = { # mutable-ok: API message payload key: value for key, value in data.items() if key != "system" } @@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation): ] ) - # Step 1: Extract all text content and images + # Step 1: Extract all text content, images, and tool calls + top_level_system_scanned: Final = ( + () + if hoisted_system_message is None or scan_only_tool_results + else self._extract_top_level_system_text(hoisted_system_message) + ) extracted: Final = tuple( self._extract_input_text_and_images( message=message, @@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation): ) for msg_idx, message in enumerate(messages) ) - scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned) + scanned: Final = ( + *top_level_system_scanned, + *(item for one_message in extracted for item in one_message.scanned), + ) texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str] images_to_check: Final = [ image for one_message in extracted for image in one_message.images ] # mutable-ok: GenericGuardrailAPIInputs takes list[str] + scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls) + tool_calls_to_check: Final = [ + item.tool_call for item in scanned_tool_calls + ] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk] + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: + # Step 2: Apply guardrail to all texts and tool calls in batch + if texts_to_check or tool_calls_to_check: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check if tools_to_check: inputs["tools"] = tools_to_check original_structured_messages: Final = structured_messages @@ -572,10 +678,16 @@ class AnthropicMessagesHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( - messages=messages, + data=data, responses=guardrailed_texts, scanned=scanned, ) + self._apply_guardrail_tool_calls_to_input( + messages=messages, + scanned_tool_calls=scanned_tool_calls, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) @@ -598,6 +710,19 @@ class AnthropicMessagesHandler(BaseTranslation): hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload return hoisted[0] if hoisted else None + @staticmethod + def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]: + content: Final = hoisted_system_message.get("content") + if isinstance(content, str): + return (ScannedText(content, SystemStringTarget()),) + if not isinstance(content, list): + return () + return tuple( + ScannedText(text_str, SystemBlockTextTarget(block_idx)) + for block_idx, block in enumerate(content) + if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str + ) + @staticmethod def _openai_system_message_to_anthropic( message: Mapping[str, object], @@ -852,9 +977,25 @@ class AnthropicMessagesHandler(BaseTranslation): for content_idx, content_item in enumerate(content) if isinstance(content_item, dict) ) + tool_use_blocks: Final = ( + () + if scan_only_tool_results + else tuple( + (content_idx, content_item) + for content_idx, content_item in enumerate(content) + if isinstance(content_item, dict) and _is_client_tool_use(content_item) + ) + ) return ExtractedInput( scanned=tuple(item for block in blocks for item in block.scanned), images=tuple(image for block in blocks for image in block.images), + tool_calls=tuple( + ScannedToolCall( + tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx), + target=ToolUseInputTarget(msg_idx, content_idx), + ) + for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks) + ), ) @classmethod @@ -940,43 +1081,48 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: Sequence[_WritableMessage], + data: dict, responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: """ - Apply guardrail responses back to input messages. + Apply guardrail responses back to the top-level system prompt and the input messages. """ + messages: Final[Sequence[_WritableMessage]] = data.get("messages") or () for item, guardrail_response in zip(scanned, responses): - target = item.target - message = messages[target.msg_idx] - content = message.get("content", None) - if content is None: - continue - - match target: - case MessageContentTarget(): - if isinstance(content, str): - message["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ContentBlockTextTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["text"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultStringTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): - if isinstance(content, list): - content[content_idx]["content"][block_idx]["text"] = ( + match item.target: + case SystemStringTarget(): + if isinstance(data.get("system"), str): + data["system"] = ( guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place ) + case SystemBlockTextTarget(block_idx=block_idx): + _write_back_system_block(data.get("system"), block_idx, guardrail_response) + case ( + MessageContentTarget() + | ContentBlockTextTarget() + | ToolResultStringTarget() + | ToolResultBlockTextTarget() as message_target + ): + _write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response) case _: - assert_never(target) + assert_never(item.target) + + @staticmethod + def _apply_guardrail_tool_calls_to_input( + messages: Sequence[_WritableMessage], + scanned_tool_calls: tuple[ScannedToolCall, ...], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + returned_tool_calls: object, + ) -> None: + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tuple(item.tool_call for item in scanned_tool_calls) + ) + for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): + if before != after: + _write_back_tool_use(messages[item.target.msg_idx], item.target, after) async def process_output_response( self, 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 9fe56f4dc65..4b47b7d8c4a 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 @@ -635,14 +635,19 @@ class TestAnthropicMessagesHandlerInputProcessing: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert guardrail.inputs is not None - assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"] + assert guardrail.inputs["texts"] == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + ] structured = guardrail.inputs["structured_messages"] assert [m["role"] for m in structured] == ["system", "user", "system"] assert structured[0]["content"] == "trusted top-level system prompt" + assert data["system"] == "trusted top-level system prompt" assert data["messages"][1]["content"] == "[MASKED]" @pytest.mark.asyncio - async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included( + async def test_bedrock_masking_slice_lines_up_when_top_level_system_is_included( self, ): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( @@ -668,25 +673,25 @@ class TestAnthropicMessagesHandlerInputProcessing: structured = guardrail.inputs["structured_messages"] bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") - assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1 + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") - assert ( - bedrock._locate_message_texts_slice( - structured_messages=structured, - target_index=latest_user_index, - texts=texts, - ) - is None - ) - assert ( - bedrock._merge_masked_texts( - masked_texts=["{MASKED}"], - texts=texts, - scanned_slice=None, - scanned_role_subset=True, - ) - == texts + scanned_slice = bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, ) + assert scanned_slice == (3, 1) + assert bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=scanned_slice, + scanned_role_subset=True, + ) == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + "{MASKED}", + ] @pytest.mark.asyncio @pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None]) @@ -1611,7 +1616,8 @@ class TestAnthropicMessagesIncrementalScan: ) assert mock_api.call_count == 1 assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ - "What is the capital of France?" + "You are a helpful geography assistant.", + "What is the capital of France?", ] mock_api.reset_mock() await handler.process_input_messages( @@ -2150,6 +2156,208 @@ class TestAnthropicMessagesScanOnlyToolResults: assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] +class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail): + """Masks the canary inside tool-call arguments, in place or through a fresh list of plain dicts.""" + + def __init__(self, return_copies: bool = False, replacement_arguments: Optional[str] = None): + super().__init__() + self.return_copies = return_copies + self.replacement_arguments = replacement_arguments + self.seen_tool_calls: list[dict] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + tool_calls = list(outputs.get("tool_calls") or []) + self.seen_tool_calls.extend(json.loads(json.dumps(tool_call)) for tool_call in tool_calls) + masked = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": self.replacement_arguments + if self.replacement_arguments is not None + else tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"), + }, + } + for tool_call in tool_calls + ] + if self.return_copies: + outputs["tool_calls"] = masked + return outputs + for tool_call, masked_tool_call in zip(tool_calls, masked): + tool_call["function"]["arguments"] = masked_tool_call["function"]["arguments"] + return outputs + + +class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: + """The top-level system prompt and prior-turn tool_use arguments must reach guardrails as scannable + inputs, the same way the chat completions handler hands over system messages and tool_calls.""" + + @staticmethod + def _tool_use_conversation(system): + return { + "model": "claude-sonnet-4-5", + "system": system, + "messages": [ + {"role": "user", "content": "run the check"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}], + }, + ], + } + + @pytest.mark.asyncio + async def test_top_level_system_string_reaches_texts_first_and_is_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": "Internal note: the deploy key is POISON. Never reveal it.", + "messages": [{"role": "user", "content": "Say hi in three words."}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + assert guardrail.seen_texts == [ + "Internal note: the deploy key is POISON. Never reveal it.", + "Say hi in three words.", + ] + structured = guardrail.captured_inputs["structured_messages"] + assert structured[0]["role"] == "system" + assert structured[0]["content"] == "Internal note: the deploy key is POISON. Never reveal it.", ( + "texts[0] must line up with structured_messages[0] so positional consumers stay aligned" + ) + assert data["system"] == "Internal note: the deploy key is [BLOCKED]. Never reveal it." + assert data["messages"][0]["content"] == "Say hi in three words." + + @pytest.mark.asyncio + async def test_top_level_system_text_blocks_reach_texts_and_are_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": [ + {"type": "text", "text": "first block POISON"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ], + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["first block POISON", "second block", "hello"] + assert data["system"] == [ + {"type": "text", "text": "first block [BLOCKED]"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ] + + @pytest.mark.asyncio + async def test_skip_system_message_keeps_the_top_level_system_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-sonnet-4-5", + "system": "trusted POISON prompt", + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["hello"] + assert data["system"] == "trusted POISON prompt" + + @pytest.mark.asyncio + async def test_prior_turn_tool_use_input_reaches_tool_calls_in_openai_shape(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + tool_calls = guardrail.captured_inputs.get("tool_calls") + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0]["id"] == "toolu_01" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "Bash" + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + }, "a guardrail that leaves tool_calls alone must leave the tool_use input alone" + + @pytest.mark.asyncio + @pytest.mark.parametrize("return_copies", [False, True]) + async def test_masked_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies) + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"] + tool_use = data["messages"][1]["content"][0] + assert tool_use == { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"}, + } + assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" + + @pytest.mark.asyncio + async def test_non_json_rewritten_arguments_keep_the_tool_use_input(self): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + + @pytest.mark.asyncio + async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.scan_only_tool_results = True + data = self._tool_use_conversation(system="trusted POISON prompt") + data["messages"][2]["content"][0]["content"] = "fetched POISON page" + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["fetched POISON page"] + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("tool_calls") is None + assert data["system"] == "trusted POISON prompt" + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][2]["content"][0]["content"] == "fetched [BLOCKED] page" + + class TestStructuredWriteBackKeepsToolResults: """A guardrail rewrite must never leave a tool_use without its tool_result (Claude Code ToolSearch, LIT-6103).""" From 49b2d71057d4cff4e3a4baad843a7db7ac35c7c2 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sun, 13 Sep 2026 08:47:47 +0000 Subject: [PATCH 069/297] test(guardrails): type the native lifecycle logging_only test double Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_litellm/integrations/test_custom_guardrail.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 8095be3c0de..bb29bfed283 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2902,9 +2902,15 @@ class _NativeLifecycleLoggingGuardrail(CustomGuardrail): guardrail_name="native-logging-guardrail", event_hook=GuardrailEventHooks.logging_only, ) - self.calls: list = [] + self.calls: list[tuple[Literal["request", "response"], list[str]]] = [] - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: self.calls.append((input_type, list(inputs.get("texts") or []))) return inputs From 6f7882db34b4d15a0dec402b6218c59afc58c012 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:54:12 -0700 Subject: [PATCH 070/297] fix(types): import CachedTokensDetails on its own line in openai.py CodeQL resolves `from openai import Omit` in litellm/types/llms/openai.py to the module itself, so every importer of a name whose definition line is in the diff is reported as an unsafe cyclic import. 76ae35dfcd edited the line that defines BaseLiteLLMOpenAIResponseObject there and got two alerts at files this PR does not touch. That line is now byte-identical to main and CachedTokensDetails arrives through a relative import isort keeps separate. --- litellm/types/llms/openai.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 9bdad700d4b..e3eac9b9205 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -82,7 +82,7 @@ from typing_extensions import ( override, ) -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, CachedTokensDetails +from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.responses.main import ( CustomToolCallOutputItem, GenericResponseOutputItem, @@ -91,6 +91,8 @@ from litellm.types.responses.main import ( OutputImageGenerationCall, ) +from .base import CachedTokensDetails + FileContent = IO[bytes] | bytes | PathLike FileTypes = ( From c0c0c9a9ebc7443b1724c22e3b6b856aac1f750d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:05:10 -0700 Subject: [PATCH 071/297] fix(sdk): keep body and proxy headers on BadRequestError mapped from a litellm_proxy 400 The generic 400 branch of the OpenAI exception mapper dropped the wire body and no branch carried the response headers, so an application calling a LiteLLM proxy through a litellm_proxy/ model could not tell a guardrail block from any other failure without walking __cause__. BadRequestError now takes headers, filled for a litellm_proxy upstream, and the generic branch passes the body. The proxy edge treats the literal "None" type and param an older proxy sends as absent and stops forwarding an upstream proxy's date and server headers. --- litellm/constants.py | 6 +- litellm/exceptions.py | 4 +- .../exception_mapping_utils.py | 16 +++++ .../common_utils/openai_error_payload.py | 6 +- .../test_exception_mapping_utils.py | 68 +++++++++++++++++++ .../common_utils/test_openai_error_payload.py | 18 +++++ .../proxy/test_common_request_processing.py | 15 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 8 files changed, 129 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..caf1aff6792 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1964,7 +1964,11 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( } ) -UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS +ORIGIN_SERVER_HEADERS: Final[frozenset[str]] = frozenset({"date", "server"}) + +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( + HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS +) # A retrieved response replays the usage of the call that created it, so pricing these # read/management routes like inference bills the same tokens twice. diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3f22a4b2dcd..3d9e26e450b 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,7 +10,7 @@ ## LiteLLM versions of the OpenAI Exception Types import enum -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -226,6 +226,7 @@ class BadRequestError(openai.BadRequestError): max_retries: int | None = None, num_retries: int | None = None, body: dict | None = None, + headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -234,6 +235,7 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 82708d412c9..fd1ba666887 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,6 +1,7 @@ import json import re import traceback +from collections.abc import Mapping from typing import Any, Final, Protocol, cast import httpx @@ -254,6 +255,15 @@ class _ProviderHTTPException(Protocol): llm_provider: str +def _litellm_proxy_response_headers( + original_exception: _ProviderHTTPException, custom_llm_provider: str +) -> Mapping[str, str] | None: + if custom_llm_provider != "litellm_proxy": + return None + headers: Final = getattr(original_exception, "headers", None) + return headers if isinstance(headers, Mapping) else None + + def _map_openai_exception( *, model: str, @@ -264,6 +274,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: + upstream_headers: Final = _litellm_proxy_response_headers(original_exception, custom_llm_provider) # custom_llm_provider is openai, make it OpenAI message = get_error_message(error_obj=original_exception) if message is None: @@ -348,6 +359,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: raise BadRequestError( @@ -357,6 +369,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif ( "Web server is returning an unknown error" in error_str @@ -404,6 +417,8 @@ def _map_openai_exception( model=model, response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif original_exception.status_code == 401: raise AuthenticationError( @@ -436,6 +451,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif original_exception.status_code == 429: raise RateLimitError( diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 89f735ee8b6..90b3c998247 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -8,6 +8,8 @@ from typing import Final from fastapi import status +_STRINGIFIED_NONE: Final = "None" + _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { status.HTTP_401_UNAUTHORIZED: "authentication_error", @@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str: """OpenAI types ``error.type`` as a required string, so an exception carrying none falls back to the type its status code stands for.""" carried: Final = attribute_of(exc, "type") - if isinstance(carried, str): + if isinstance(carried, str) and carried != _STRINGIFIED_NONE: return carried mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: @@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None: """OpenAI types ``error.param`` as nullable, so an exception carrying none serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") - return carried if isinstance(carried, str) else None + return carried if isinstance(carried, str) and carried != _STRINGIFIED_NONE else None diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 42d3df76902..ea6ac17ad45 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1409,3 +1409,71 @@ def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): exception_headers = _get_response_headers(original_exception=exc_info.value) assert exception_headers is not None assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7 + + +_GUARDRAIL_BLOCK_ERROR = { + "message": "Content blocked: secret_project_codename pattern detected", + "param": "None", + "code": "400", + "provider_specific_fields": { + "error": "Content blocked: secret_project_codename pattern detected", + "pattern": "secret_project_codename", + "guardrail_name": "block-secret-project", + "guardrail_mode": "pre_call", + }, +} + + +def _openai_handler_error(error_type: str, headers: dict[str, str]) -> OpenAIError: + """What litellm/llms/openai/openai.py raises after the openai SDK rejects a 400: + the SDK's str() carries the wire body, and the handler copies headers and body over.""" + wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type} + wire = httpx.Response( + status_code=400, + headers=headers, + json={"error": wire_error}, + request=httpx.Request("POST", "http://localhost:4000/v1/chat/completions"), + ) + return OpenAIError( + status_code=400, + message=f"Error code: 400 - {{'error': {wire_error}}}", + headers=wire.headers, + body=wire_error, + ) + + +@pytest.mark.parametrize("error_type", ["None", "invalid_request_error"]) +def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): + """An SDK caller behind a proxy tells a guardrail block from any other 400 by the body's + provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError + must carry both whichever error.type the proxy version on the other end emits.""" + proxy_headers = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} + + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error(error_type, proxy_headers), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" + assert exc_info.value.body["type"] == error_type + assert proxy_headers.items() <= exc_info.value.headers.items() + + +def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): + """A vendor's own response headers stay on e.response the way every other mapped provider + error keeps them; only a LiteLLM proxy upstream puts headers on e.headers.""" + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="gpt-5.4-mini", + original_exception=_openai_handler_error("vendor_specific_error", {"openai-organization": "org-1"}), + custom_llm_provider="openai", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body["type"] == "vendor_specific_error" + assert exc_info.value.headers is None diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 8b653ddfb71..db9fe3a2253 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -143,3 +143,21 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): exc = HTTPException(status_code=403, detail="blocked by policy") assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" + + +def test_the_stringified_none_an_older_upstream_proxy_sent_is_treated_as_absent(): + """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on + the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal + is the exact bug this module exists to stop.""" + from litellm.exceptions import BadRequestError + + carried = BadRequestError( + message="Content blocked", + model="claude-haiku-4-5", + llm_provider="litellm_proxy", + body={"message": "Content blocked", "type": "None", "param": "None", "code": "400"}, + ) + + assert carried.type == "None" + assert openai_error_type(carried, 400) == "invalid_request_error" + assert openai_error_param(carried) is None diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 69e89d1c604..3cb58c44354 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3985,6 +3985,21 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-custom-safe"] == "1" assert proxy_exc.headers["x-request-id"] == "abc-123" + async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): + """A proxy fronting another LiteLLM proxy gets the upstream's date and server + on the mapped exception; forwarding them would duplicate the Date header + uvicorn adds to every response and leak the upstream server identity.""" + exc = litellm.BadRequestError( + message="Content blocked", + llm_provider="litellm_proxy", + model="claude-haiku-4-5", + headers={"date": "Sun, 13 Sep 2026 08:43:51 GMT", "server": "uvicorn", "x-request-id": "abc-123"}, + ) + proxy_exc = await self._invoke(exc) + assert "date" not in proxy_exc.headers + assert "server" not in proxy_exc.headers + assert proxy_exc.headers["x-request-id"] == "abc-123" + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 76b26e41abfe6e72dd846e7272dcb45069b98b73 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:07:53 +0000 Subject: [PATCH 072/297] fix(router): cool down team deployments on 429 when a sibling serves the same public model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 10 ++++ litellm/router_utils/cooldown_handlers.py | 5 +- .../router_utils/test_cooldown_handlers.py | 50 +++++++++++++++++++ 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..4929b17f7fc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1553,6 +1553,16 @@ class Router: return False return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1 + def team_model_has_alternatives(self, deployment_id: str) -> bool: + deployment: Final = self.get_deployment(model_id=deployment_id) + if deployment is None: + return False + team_id: Final = deployment.model_info.team_id + public_model_name: Final = deployment.model_info.team_public_model_name + if team_id is None or public_model_name is None: + return False + return len(self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()) > 1 + _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None: diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index f722b6fd20c..027f0a9ca05 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -343,8 +343,9 @@ def _should_cooldown_deployment( model_group: Final = litellm_router_instance.get_model_group(id=deployment) is_single_deployment_model_group = False if model_group is not None and len(model_group) == 1: - is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives( - requested_model_group + is_single_deployment_model_group = not ( + litellm_router_instance.routing_group_has_alternatives(requested_model_group) + or litellm_router_instance.team_model_has_alternatives(deployment) ) ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 7ee0ed3701b..6f66bb863cc 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -437,3 +437,53 @@ class TestRoutingGroupCooldownAlternatives: ) is False ) + + +class TestTeamModelCooldownAlternatives: + def _router(self, team_deployments: int): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": f"model_name_team-1_{i}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": f"team-deploy-{i}", + "team_id": "team-1", + "team_public_model_name": "team-gpt-4o-mini", + }, + } + for i in range(team_deployments) + ] + ) + + def test_429_on_team_deployment_with_sibling_cools_down(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is True + ) + + def test_429_on_only_team_deployment_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=1) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + ) From 073d4fe2b01500526829523f6596a60428fbead9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:11:06 -0700 Subject: [PATCH 073/297] fix(responses): route mid-stream error events through exception_type so content_policy_fallbacks fire Mid-stream error events on the streaming Responses API were all raised as APIError, so a content_policy_violation event never matched the router's content-policy fallback dispatch and the client got the raw error instead of the fallback model's answer. Map each error event's code and status through the existing exception_type mapping, matching the non-streaming path, and unwrap the typed ContentPolicyViolationError and ContextWindowExceededError so the router routes them to the configured content_policy_fallbacks and context_window_fallbacks. --- litellm/responses/streaming_iterator.py | 48 ++++-- litellm/router.py | 10 +- .../test_streaming_iterator_error_events.py | 157 ++++++++++++++++-- tests/test_litellm/test_router.py | 105 ++++++++++++ 4 files changed, 288 insertions(+), 32 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 40ff88fc557..b3426fbbfef 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -221,6 +222,13 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None ) +def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: + if isinstance(mapped_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError)): + return True + status_code: Final = getattr(mapped_exception, "status_code", None) + return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -521,15 +529,8 @@ class BaseResponsesAPIStreamingIterator: getattr(self.completed_response, "response", None) if self.completed_response else None ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None - error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) - exception: Final = litellm.APIError( - status_code=_status_code_for_error_fields(error_type, error_code), - message=error_message, - llm_provider=self.custom_llm_provider or "", - model=self.model or "", - ) - self._handle_failure(exception) + self._handle_failure(self._map_error_event_exception(error_info)) def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: @@ -551,6 +552,26 @@ class BaseResponsesAPIStreamingIterator: self.logging_obj._response_cost_calculator(result=response_obj) or 0.0 ) + def _map_error_event_exception(self, error_obj: object) -> Exception: + error_message, error_type, error_code = _error_event_fields(error_obj) + status_code: Final = _status_code_for_error_fields(error_type, error_code) + error_body: Final = {"message": error_message, "type": error_type, "code": error_code} + provider_exception: Final = BaseLLMException( + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {error_body}}}", + body=error_body, + ) + try: + return litellm.exception_type( + model=self.model or "", + custom_llm_provider=self.custom_llm_provider or "", + original_exception=provider_exception, + completion_kwargs={}, + extra_kwargs={}, + ) + except Exception as mapped_exception: + return mapped_exception + def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) if chunk_type not in ("error", "response.failed"): @@ -562,15 +583,8 @@ class BaseResponsesAPIStreamingIterator: else getattr(result, "error", None) ) - error_message, error_type, error_code = _error_event_fields(error_obj) - status_code: Final = _status_code_for_error_fields(error_type, error_code) - mapped_exception: Final = litellm.APIError( - status_code=status_code, - message=error_message, - llm_provider=self.custom_llm_provider or "", - model=self.model or "", - ) - if 400 <= status_code < 500 and status_code != 429: + mapped_exception: Final = self._map_error_event_exception(error_obj) + if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception raise MidStreamFallbackError( message=str(mapped_exception), diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..79854a11150 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3271,8 +3271,16 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + fallback_trigger: Final[Exception] = ( + e.original_exception + if isinstance( + e.original_exception, + (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), + ) + else e + ) fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e, + e=fallback_trigger, disable_fallbacks=False, fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index ad74861c096..73afbb5e63a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -1,9 +1,14 @@ """ Regression: in-stream error events (type="error", type="response.failed") must raise instead of being returned as benign chunks, mirroring chat streaming -semantics (_handle_stream_fallback_error): non-retriable 4xx (except 429) -raise litellm.APIError directly; 429 and 5xx are wrapped in -MidStreamFallbackError so the Router's mid-stream fallback machinery fires. +semantics (_handle_stream_fallback_error). The event's code, type and status go +through litellm.exception_type, so each event raises the same typed exception +the non-streaming path raises for that provider error: non-retriable 4xx +(except 429) raise that typed exception directly, while 429, 5xx, +ContentPolicyViolationError and ContextWindowExceededError are wrapped in +MidStreamFallbackError so the Router's mid-stream fallback machinery fires and +its content_policy_fallbacks / context_window_fallbacks dispatch sees the +trigger it matches on. Status mapping must consider both the OpenAI error `type` (e.g. "invalid_request_error") and `code` (e.g. "invalid_prompt", @@ -66,12 +71,12 @@ def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback( with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 500 - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.InternalServerError) assert exc_info.value.original_exception.status_code == 500 def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fallback(): - """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped APIError.""" + """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped RateLimitError.""" iterator = _make_iterator() chunk = _make_error_chunk("tokens", "rate_limit_exceeded", "Too many requests") with pytest.raises(MidStreamFallbackError) as exc_info: @@ -79,15 +84,15 @@ def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fall assert exc_info.value.status_code == 429 assert exc_info.value.generated_content == "" assert exc_info.value.is_pre_first_chunk is True - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) assert exc_info.value.original_exception.status_code == 429 def test_maybe_raise_for_error_event_maps_invalid_request_type_to_400(): - """Client errors classified via the `type` field must raise APIError directly (no fallback).""" + """Client errors classified via the `type` field must raise BadRequestError directly (no fallback).""" iterator = _make_iterator() chunk = _make_error_chunk("invalid_request_error", "invalid_prompt", "bad request") - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) @@ -99,12 +104,84 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): chunk = Mock() chunk.type = "error" chunk.error = {"code": "context_length_exceeded", "message": "too long"} - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) +def test_maybe_raise_for_error_event_wraps_context_window_exceeded_for_context_window_fallbacks(): + """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming + path and is wrapped so the Router's context_window_fallbacks dispatch fires mid-stream.""" + iterator = _make_iterator() + chunk = _make_error_chunk( + "invalid_request_error", + "context_length_exceeded", + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContextWindowExceededError) + assert exc_info.value.status_code == 400 + + +CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) +def test_maybe_raise_for_error_event_wraps_content_policy_violation_for_content_policy_fallbacks( + custom_llm_provider: str, +): + """Regression: a content_policy_violation error event used to raise a bare APIError, so the Router's + content_policy_fallbacks never fired. It must map to ContentPolicyViolationError (the same exception the + non-streaming path raises) and be wrapped so the Router's mid-stream fallback catches it.""" + iterator = _make_iterator() + iterator.custom_llm_provider = custom_llm_provider + chunk = _make_error_chunk("invalid_request_error", "content_policy_violation", CONTENT_POLICY_MESSAGE) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + assert exc_info.value.original_exception.status_code == 400 + assert exc_info.value.status_code == 400 + assert exc_info.value.is_pre_first_chunk is True + assert CONTENT_POLICY_MESSAGE in str(exc_info.value.original_exception) + + +def test_maybe_raise_for_response_failed_event_wraps_content_policy_violation(): + iterator = _make_iterator() + chunk = _make_failed_chunk( + {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE} + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + + +@pytest.mark.parametrize( + "error_type,error_code,expected_exception", + [ + ("invalid_request_error", "content_policy_violation", litellm.ContentPolicyViolationError), + ("tokens", "rate_limit_exceeded", litellm.RateLimitError), + ("invalid_request_error", "insufficient_quota", litellm.RateLimitError), + ("server_error", "internal_error", litellm.InternalServerError), + ("invalid_request_error", "invalid_prompt", litellm.BadRequestError), + ("invalid_request_error", "model_not_found", litellm.NotFoundError), + ("server_error", "vector_store_timeout", litellm.Timeout), + ], +) +def test_error_event_raises_the_same_typed_exception_as_the_non_streaming_path( + error_type: str, error_code: str, expected_exception: type[Exception] +): + iterator = _make_iterator() + chunk = _make_error_chunk(error_type, error_code, "provider message") + with pytest.raises((MidStreamFallbackError, expected_exception)) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + raised = exc_info.value + typed_exception = raised.original_exception if isinstance(raised, MidStreamFallbackError) else raised + assert type(typed_exception) is expected_exception + assert "provider message" in str(typed_exception) + + def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): """OpenAI returns HTTP 429 for insufficient_quota; it must not map to 400 even though its type is invalid_request_error-adjacent, and it must be wrapped for fallback.""" @@ -113,6 +190,7 @@ def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) def test_maybe_raise_for_error_event_passes_through_normal_chunk(): @@ -186,10 +264,43 @@ async def test_async_iterator_raises_mid_stream_fallback_on_rate_limit_error_eve assert exc_info.value.status_code == 429 assert exc_info.value.is_pre_first_chunk is True assert exc_info.value.generated_content == "" - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) assert exc_info.value.original_exception.status_code == 429 +@pytest.mark.asyncio +async def test_async_iterator_content_policy_violation_after_first_chunk_carries_generated_content(): + """The customer's case: text streams, then the provider halts the stream with a + content_policy_violation error event. The iterator must surface ContentPolicyViolationError + inside MidStreamFallbackError, together with the text already streamed.""" + iterator = _make_async_iterator_with_events( + [ + {"type": "response.output_text.delta", "delta": "partial "}, + { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": CONTENT_POLICY_MESSAGE, + }, + }, + ] + ) + + chunks = [] + + async def _drain(): + async for chunk in iterator: + chunks.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _drain() + assert len(chunks) == 1 + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + assert exc_info.value.is_pre_first_chunk is False + assert exc_info.value.generated_content == "partial " + + @pytest.mark.asyncio async def test_async_iterator_error_after_first_chunk_carries_generated_content(): """An error after streamed output must expose the accumulated text so the router's @@ -265,7 +376,7 @@ def test_handle_logging_failed_response_maps_rate_limit_to_429(): ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] - assert isinstance(logged_exception, litellm.APIError) + assert isinstance(logged_exception, litellm.RateLimitError) assert logged_exception.status_code == 429 assert "throttled" in str(logged_exception) @@ -282,10 +393,28 @@ def test_handle_logging_failed_response_maps_type_field_to_400(): ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] - assert isinstance(logged_exception, litellm.APIError) + assert isinstance(logged_exception, litellm.BadRequestError) assert logged_exception.status_code == 400 +def test_handle_logging_failed_response_logs_content_policy_violation(): + """Failure logging must record the same typed exception the stream raises, so logging + integrations see a content policy violation instead of a generic APIError.""" + iterator = _make_iterator() + iterator.completed_response = _make_failed_chunk( + {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE} + ) + with ( + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function") as mock_run_async, + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), + ): + iterator._handle_logging_failed_response() + logged_exception = mock_run_async.call_args.kwargs["exception"] + assert isinstance(logged_exception, litellm.ContentPolicyViolationError) + assert logged_exception.status_code == 400 + assert CONTENT_POLICY_MESSAGE in str(logged_exception) + + def test_handle_logging_failed_response_records_usage_and_cost(): """Usage on a response.failed event must reach failure spend accounting via combined_usage_object.""" iterator = _make_iterator() @@ -357,7 +486,7 @@ def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event(): for _ in iterator: pass assert exc_info.value.status_code == 429 - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) def test_every_openai_sdk_response_error_code_has_explicit_status_mapping(): @@ -413,7 +542,7 @@ def test_maybe_raise_for_response_failed_event_maps_image_code_to_400(): chunk = Mock() chunk.type = "response.failed" chunk.response = mock_response_obj - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..d74074dd149 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3654,6 +3654,111 @@ async def test_aresponses_streaming_iterator_fallback(): assert call_kwargs["disable_fallbacks"] is False +@pytest.mark.asyncio +async def test_aresponses_streaming_content_policy_error_event_routes_to_content_policy_fallback(): + """Regression: a mid-stream content_policy_violation error event never reached + content_policy_fallbacks. The iterator raised a bare APIError the wrapper does not + catch, and even once wrapped, the MidStreamFallbackError envelope was handed to the + fallback dispatch, whose isinstance branch on ContentPolicyViolationError never matched. + The stream below is the customer's shape: a raw OpenAI error event with code + content_policy_violation, transformed by the real OpenAI config, and the router must + call the content_policy_fallbacks target, not the general fallbacks one.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + router = litellm.Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k1"}}, + { + "model_name": "content-fallback", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k2"}, + }, + {"model_name": "general-fallback", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["general-fallback"]}], + content_policy_fallbacks=[{"primary": ["content-fallback"]}], + ) + error_event = { + "type": "error", + "sequence_number": 2, + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": "This content was flagged for possible cybersecurity risk. The response was halted mid-stream.", + "param": None, + }, + } + + async def aiter_bytes(): + yield f"data: {json.dumps(error_event)}\n\n".encode() + + raw_response = MagicMock() + raw_response.headers = {} + raw_response.aiter_bytes = aiter_bytes + logging_obj = MagicMock(spec=LiteLLMLogging) + logging_obj.model_call_details = {"litellm_params": {}} + logging_obj.completion_start_time = None + source = ResponsesAPIStreamingIterator( + response=raw_response, + model="gpt-5.4", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + fallback_chunks = [MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed")] + fallback_call = AsyncMock(return_value=_AsyncList(fallback_chunks)) + + wrapped = await router._aresponses_streaming_iterator( + response=source, + initial_kwargs={ + "model": "primary", + "stream": True, + "input": "Hi", + "original_generic_function": fallback_call, + }, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == fallback_chunks + fallback_call.assert_awaited_once() + assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): + """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the + ContentPolicyViolationError carried inside MidStreamFallbackError, not the envelope.""" + router = _make_router_with_fallback("openai/gpt-5.4", "openai/gpt-5-mini") + content_policy_error = litellm.ContentPolicyViolationError( + message="flagged mid-stream", llm_provider="openai", model="openai/gpt-5.4" + ) + src = _make_responses_iterator( + chunks=[MagicMock(type="response.created")], + error=MidStreamFallbackError( + message=str(content_policy_error), + model="openai/gpt-5.4", + llm_provider="openai", + original_exception=content_policy_error, + is_pre_first_chunk=True, + ), + model="openai/gpt-5.4", + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AsyncList([MagicMock(type="response.completed")])), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={"model": "openai/gpt-5.4", "stream": True, "input": "Hi"}, + ) + [chunk async for chunk in wrapped] + + mock_fallback_utils.assert_awaited_once() + assert mock_fallback_utils.await_args.kwargs["e"] is content_policy_error + + @pytest.mark.asyncio @pytest.mark.parametrize( "fallback_headers", From e732a484f6b427d6513bc39885fdc8485ba29e7c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:11:59 -0700 Subject: [PATCH 074/297] fix(responses): hoist Codex additional_tools input items into the chat bridge tools --- .../responses/transformation.py | 66 +++--------- litellm/responses/additional_tools.py | 65 +++++++++++ .../custom_tools.py | 28 +++-- .../handler.py | 22 ++-- .../transformation.py | 22 ++-- .../test_handler.py | 102 ++++++++++++++++++ .../test_litellm_completion_responses.py | 91 ++++++++++++++++ .../responses/test_additional_tools.py | 48 +++++++++ .../responses/test_custom_tool_call.py | 18 ++++ 9 files changed, 387 insertions(+), 75 deletions(-) create mode 100644 litellm/responses/additional_tools.py create mode 100644 tests/test_litellm/responses/test_additional_tools.py diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 53a3e634adf..95375399033 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -17,7 +17,7 @@ BaseAWSLLM._sign_request after the request body is finalized. import json from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict import httpx from typing_extensions import ReadOnly, TypedDict @@ -32,6 +32,7 @@ from litellm.llms.bedrock_mantle.common_utils import ( BedrockMantleAuthMixin, ) from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( ResponseInputParam, @@ -58,8 +59,6 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset( _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"}) -_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" - _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message" _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction" _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call" @@ -233,62 +232,29 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: - remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input) - normalized_input: Final = self._normalize_codex_input_items(remaining_input) - request_params: Final = ( - { - **response_api_optional_request_params, - "tools": [ - *(response_api_optional_request_params.get("tools") or []), - *hoisted_tools, - ], - } - if hoisted_tools - else response_api_optional_request_params + params: Final = cast( # cast-ok: the base signature leaves the params dict untyped + "ResponsesAPIOptionalRequestParams", response_api_optional_request_params ) + hoisted: Final = hoist_additional_tools(input, params.get("tools")) + normalized_input: Final = self._normalize_codex_input_items(hoisted.input) return super().transform_responses_api_request( model=model, input=normalized_input, - response_api_optional_request_params=request_params, + response_api_optional_request_params=self._params_with_hoisted_tools(params, hoisted), litellm_params=litellm_params, headers=headers, ) - @staticmethod - def _is_codex_additional_tools_item(item: Any) -> bool: - return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE - - @staticmethod - def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]": - tools: Final = item.get("tools") - return tools if isinstance(tools, list) else [] - @classmethod - def _hoist_codex_additional_tools( - cls, - input: "str | ResponseInputParam", - ) -> "tuple[str | ResponseInputParam, list[Any]]": - """Codex's "responses lite" wire mode ships tool definitions inside - `input` as {"type": "additional_tools", "role": "developer", - "tools": [...]} items. api.openai.com accepts that item type; Mantle - rejects the whole request with 400 "Invalid 'input': value did not - match any expected variant" but accepts the same tools at the top - level, so move them there and strip the items from `input`. - """ - if not isinstance(input, list): - return input, [] - additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)] - if not additional_tools_items: - return input, [] - remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)] - hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)] - verbose_logger.debug( - "Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) " - "into the top-level tools param (Mantle rejects that input item type).", - len(hoisted_tools), - len(additional_tools_items), - ) - return remaining_input, cls._filter_unsupported_tools(hoisted_tools) + def _params_with_hoisted_tools( + cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools + ) -> dict[str, object]: + if not hoisted.hoisted: + return dict(params) + supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools)) + if supported_tools: + return {**params, "tools": supported_tools} + return {key: value for key, value in params.items() if key != "tools"} @staticmethod def _agent_message_text(item: "Mapping[str, object]") -> str: diff --git a/litellm/responses/additional_tools.py b/litellm/responses/additional_tools.py new file mode 100644 index 00000000000..ea0d7af350c --- /dev/null +++ b/litellm/responses/additional_tools.py @@ -0,0 +1,65 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools + +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_logger +from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam + +ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" + + +class _InputItemType(BaseModel): + type: str = "" + + +class _AdditionalToolsItem(BaseModel): + tools: tuple[dict[str, object], ...] = () + + +@dataclass(frozen=True, slots=True) +class HoistedAdditionalTools: + input: str | ResponseInputParam + tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] + hoisted: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] + + +def _is_additional_tools_item(item: object) -> bool: + try: + return _InputItemType.model_validate(item).type == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE + except ValidationError: + return False + + +def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: + try: + parsed: Final = _AdditionalToolsItem.model_validate(item) + except ValidationError: + return () + return tuple( + cast( + "ALL_RESPONSES_API_TOOL_PARAMS", tool + ) # cast-ok: nested tools carry the same raw tool JSON as top-level tools + for tool in parsed.tools + ) + + +def hoist_additional_tools( + input: str | ResponseInputParam, + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, +) -> HoistedAdditionalTools: + existing: Final = tuple(tools or ()) + if isinstance(input, str): + return HoistedAdditionalTools(input=input, tools=existing, hoisted=()) + items: Final = tuple(item for item in input if _is_additional_tools_item(item)) + if not items: + return HoistedAdditionalTools(input=input, tools=existing, hoisted=()) + hoisted: Final = tuple(tool for item in items for tool in _tools_of_item(item)) + verbose_logger.debug( + "Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) into the top-level tools param.", + len(hoisted), + len(items), + ) + remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)] + return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index 4aa489d9e50..038964055c3 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -39,15 +39,27 @@ def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str: return f"{prefix}_{tool_id}" +class _ToolNameFields(BaseModel): + type: str = "" + name: str = "" + tools: tuple[object, ...] = () + + +def _custom_tool_names_of(tool: object) -> tuple[str, ...]: + try: + parsed: Final = _ToolNameFields.model_validate(tool) + except ValidationError: + return () + if parsed.type == "custom": + return (parsed.name,) if parsed.name else () + if parsed.type != "namespace": + return () + return tuple(name for nested_tool in parsed.tools for name in _custom_tool_names_of(nested_tool)) + + def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]: - """Extract names of tools originally defined as ``type: "custom"``.""" - if not tools: - return set() - names: Final[set[str]] = set() - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "custom" and "name" in tool: - names.add(tool["name"]) - return names + """Extract names of tools defined as ``type: "custom"``, at the top level or inside a ``namespace`` tool.""" + return {name for tool in tools or () for name in _custom_tool_names_of(tool)} def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index a0e8cd278e6..505b5b09433 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -6,6 +6,7 @@ from collections.abc import Coroutine, Mapping from typing import Final import litellm +from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) @@ -37,11 +38,16 @@ class LiteLLMCompletionTransformationHandler: | BaseResponsesAPIStreamingIterator | Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator] ): + hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools")) + bridged_input: Final = hoisted.input + bridged_request: Final[ResponsesAPIOptionalRequestParams] = ( + {**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request + ) litellm_completion_request: Final[dict] = ( LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, - input=input, - responses_api_request=responses_api_request, + input=bridged_input, + responses_api_request=bridged_request, custom_llm_provider=custom_llm_provider, stream=stream, extra_headers=extra_headers, @@ -52,8 +58,8 @@ class LiteLLMCompletionTransformationHandler: if _is_async: return self.async_response_api_handler( litellm_completion_request=litellm_completion_request, - request_input=input, - responses_api_request=responses_api_request, + request_input=bridged_input, + responses_api_request=bridged_request, **kwargs, ) @@ -70,8 +76,8 @@ class LiteLLMCompletionTransformationHandler: responses_api_response: Final[ResponsesAPIResponse] = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( chat_completion_response=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, + request_input=bridged_input, + responses_api_request=bridged_request, ) ) @@ -81,8 +87,8 @@ class LiteLLMCompletionTransformationHandler: return LiteLLMCompletionStreamingIterator( model=model, litellm_custom_stream_wrapper=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, + request_input=bridged_input, + responses_api_request=bridged_request, custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e8aacac9e67..27756b405ec 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1890,9 +1890,21 @@ class LiteLLMCompletionResponsesConfig: namespace_tool: NamespaceTool, nested: bool, ) -> ChatCompletionToolParam | None: - if nested and namespace_tool.get("type") != "function": + tool_type: Final = namespace_tool.get("type") + if nested and tool_type not in ("function", "custom"): return None + raw_description: Final = str(namespace_tool.get("description") or "") + description: Final = ( + f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" + if nested and namespace_description and raw_description + else namespace_description + if nested and namespace_description + else raw_description + ) + if nested and tool_type == "custom": + return convert_custom_tool_to_function_tool({**namespace_tool, "description": description}) + raw_parameters: Final = namespace_tool.get("parameters") parameters: Final = ( MappingProxyType(raw_parameters) if isinstance(raw_parameters, Mapping) else MappingProxyType({}) @@ -1901,14 +1913,6 @@ class LiteLLMCompletionResponsesConfig: parameters if parameters and "type" in parameters else MappingProxyType({**parameters, "type": "object"}) ) tool_name: Final = str(namespace_tool.get("name") or "") - raw_description: Final = str(namespace_tool.get("description") or "") - description: Final = ( - f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" - if nested and namespace_description and raw_description - else namespace_description - if nested and namespace_description - else raw_description - ) chat_tool_name: Final = f"{namespace}__{tool_name}" if nested else tool_name function: Final = ChatCompletionToolParamFunctionChunk( name=chat_tool_name, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py index 2cfec6a1844..b78dabbfe48 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py @@ -68,3 +68,105 @@ async def test_async_fallback_tags_skip_responses_api_bridge(): await coro assert captured.get("_skip_responses_api_bridge") is True + + +_CODEX_ADDITIONAL_TOOLS_ITEM = { + "type": "additional_tools", + "id": "at_codex", + "role": "developer", + "tools": [ + { + "type": "namespace", + "name": "functions", + "description": "", + "tools": [ + { + "type": "custom", + "name": "exec", + "description": "Runs a shell command.", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + }, + { + "type": "function", + "name": "wait", + "description": "Waits for a background command.", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}}, + }, + ], + } + ], +} +_CODEX_INPUT = [_CODEX_ADDITIONAL_TOOLS_ITEM, {"type": "message", "role": "user", "content": "Run ls"}] + + +def test_sync_fallback_hoists_additional_tools_input_items_into_chat_tools(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + def fake_completion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.completion", fake_completion): # test-quality-ok: no DI seam; the file stubs this same boundary + with pytest.raises(_StopForwarding): + handler.response_api_handler( + model="bedrock/us.openai.gpt-5.6", + input=_CODEX_INPUT, + responses_api_request={}, + custom_llm_provider="bedrock", + _is_async=False, + ) + + assert [message["role"] for message in captured["messages"]] == ["user"] + functions_by_name = {tool["function"]["name"]: tool["function"] for tool in captured["tools"]} + assert set(functions_by_name) == {"exec", "functions__wait"} + assert set(functions_by_name["exec"]["parameters"]["properties"]) == {"content"} + + +@pytest.mark.asyncio +async def test_async_fallback_returns_hoisted_nested_custom_tool_call_as_custom_tool_call(): + from litellm.responses.litellm_completion_transformation.transformation import TOOL_CALLS_CACHE + from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse + + handler = LiteLLMCompletionTransformationHandler() + tool_call_id = "call_exec_hoisted" + + async def fake_acompletion(**kwargs): + return ModelResponse( + id="chatcmpl-exec", + created=1, + model="us.openai.gpt-5.6", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id=tool_call_id, + type="function", + function=Function(name="exec", arguments='{"content": "ls"}'), + ) + ], + ), + ) + ], + ) + + try: + with patch("litellm.acompletion", fake_acompletion): # test-quality-ok: no DI seam; file stubs this boundary + response = await handler.response_api_handler( + model="bedrock/us.openai.gpt-5.6", + input=_CODEX_INPUT, + responses_api_request={}, + custom_llm_provider="bedrock", + _is_async=True, + ) + finally: + TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) + + tool_calls = [(item.type, item.name, item.input) for item in response.output if item.type == "custom_tool_call"] + assert tool_calls == [("custom_tool_call", "exec", "ls")] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 3c78bbf79d7..1f497597a11 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2506,6 +2506,7 @@ class TestToolTransformation: "tools": [ "ignored", {"type": "namespace", "name": "ignored"}, + {"type": "web_search", "name": "ignored"}, { "type": "function", "name": "spawn_agent", @@ -2527,6 +2528,36 @@ class TestToolTransformation: "type": "object", } + def test_transform_nested_namespace_custom_tool_becomes_a_content_function_under_its_short_name(self): + namespace_tool = { + "type": "namespace", + "name": "functions", + "description": "Codex shell tools.", + "tools": [ + { + "type": "custom", + "name": "exec", + "description": "Runs a shell command.", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + }, + ], + } + + result_tools, _ = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + ) + + assert len(result_tools) == 1 + function = result_tools[0]["function"] + assert function["name"] == "exec" + assert function["description"].startswith("Codex shell tools.") + assert "Runs a shell command." in function["description"] + assert "start: /.+/" in function["description"] + assert function["parameters"]["required"] == ["content"] + assert function["parameters"]["properties"]["content"]["type"] == "string" + @pytest.mark.parametrize( "model, custom_llm_provider", [ @@ -3786,6 +3817,66 @@ class TestEnsureOutputItemContentPartAdded: assert added.item.name == "spawn_agent" assert added.item.namespace == "collaboration" + def test_streaming_nested_custom_tool_call_comes_back_as_custom_tool_call(self): + from litellm.responses.litellm_completion_transformation.custom_tools import extract_custom_tool_names + + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "functions", + "tools": [ + { + "type": "custom", + "name": "exec", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + } + ], + } + ] + } + iterator._custom_tool_names = extract_custom_tool_names(iterator.responses_api_request.get("tools")) + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + + iterator._queue_tool_call_delta_events( + [{"index": 0, "id": "call_exec", "function": {"name": "exec", "arguments": '{"content":"ls"}'}}] + ) + iterator._queue_final_tool_call_done_events( + ModelResponse( + id="chatcmpl-exec", + created=1, + model="us.openai.gpt-5.6", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_exec", + type="function", + function=Function(name="exec", arguments='{"content":"ls"}'), + ) + ], + ), + ) + ], + ) + ) + + added = iterator._pending_tool_events[0] + assert added.item.type == "custom_tool_call" + assert added.item.name == "exec" + done = iterator._pending_tool_events[-1] + assert done.item.type == "custom_tool_call" + assert done.item.input == "ls" + def test_streaming_unqualified_namespace_tool_calls_restore_namespace(self): """A unique nested tool name without the namespace still maps back.""" iterator = self._make_iterator() diff --git a/tests/test_litellm/responses/test_additional_tools.py b/tests/test_litellm/responses/test_additional_tools.py new file mode 100644 index 00000000000..bef3b27eacd --- /dev/null +++ b/tests/test_litellm/responses/test_additional_tools.py @@ -0,0 +1,48 @@ +from litellm.responses.additional_tools import hoist_additional_tools + +_EXEC_TOOL = {"type": "custom", "name": "exec", "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}} +_WAIT_TOOL = {"type": "function", "name": "wait", "parameters": {"type": "object", "properties": {}}} +_TOP_LEVEL_TOOL = {"type": "function", "name": "top_level", "parameters": {"type": "object", "properties": {}}} +_USER_MESSAGE = {"type": "message", "role": "user", "content": "Run ls"} + + +def test_string_input_passes_through_with_existing_tools(): + hoisted = hoist_additional_tools("hello", [_TOP_LEVEL_TOOL]) + + assert hoisted.input == "hello" + assert hoisted.tools == (_TOP_LEVEL_TOOL,) + assert hoisted.hoisted == () + + +def test_input_without_additional_tools_items_is_returned_untouched(): + request_input = [_USER_MESSAGE] + + hoisted = hoist_additional_tools(request_input, None) + + assert hoisted.input is request_input + assert hoisted.tools == () + assert hoisted.hoisted == () + + +def test_additional_tools_items_are_stripped_and_appended_after_top_level_tools_in_item_order(): + request_input = [ + {"type": "additional_tools", "id": "at_1", "role": "developer", "tools": [_EXEC_TOOL]}, + _USER_MESSAGE, + {"type": "additional_tools", "id": "at_2", "role": "developer", "tools": [_WAIT_TOOL]}, + ] + + hoisted = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL]) + + assert hoisted.input == [_USER_MESSAGE] + assert hoisted.tools == (_TOP_LEVEL_TOOL, _EXEC_TOOL, _WAIT_TOOL) + assert hoisted.hoisted == (_EXEC_TOOL, _WAIT_TOOL) + + +def test_additional_tools_item_without_a_tools_list_is_stripped_and_contributes_nothing(): + request_input = [{"type": "additional_tools", "id": "at_1", "role": "developer", "tools": "exec"}, _USER_MESSAGE] + + hoisted = hoist_additional_tools(request_input, None) + + assert hoisted.input == [_USER_MESSAGE] + assert hoisted.tools == () + assert hoisted.hoisted == () diff --git a/tests/test_litellm/responses/test_custom_tool_call.py b/tests/test_litellm/responses/test_custom_tool_call.py index e80301c3b2f..2ed71ee3ecf 100644 --- a/tests/test_litellm/responses/test_custom_tool_call.py +++ b/tests/test_litellm/responses/test_custom_tool_call.py @@ -55,6 +55,24 @@ class TestCustomToolUtilities: names = extract_custom_tool_names(tools) assert names == set() + def test_extract_custom_tool_names_walks_namespace_tools(self): + tools = [ + {"type": "function", "name": "regular_tool"}, + { + "type": "namespace", + "name": "functions", + "tools": [ + {"type": "custom", "name": "exec"}, + {"type": "function", "name": "wait"}, + "ignored", + ], + }, + {"type": "namespace", "name": "empty", "tools": "not-a-list"}, + ] + + names = extract_custom_tool_names(tools) + assert names == {"exec"} + def test_extract_custom_tool_names_none(self): """Test extraction with None input.""" names = extract_custom_tool_names(None) From f7e9277032651786d493c72225e571d28f40c136 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:13:17 -0700 Subject: [PATCH 075/297] fix(guardrails): validate rewritten tool_use arguments with a typed adapter --- .../chat/guardrail_translation/handler.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index eb16278c9bd..b438d168b52 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -20,6 +20,7 @@ from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger @@ -188,7 +189,7 @@ def _is_client_tool_use(block: Mapping[str, object]) -> bool: block.get("type") == "tool_use" and isinstance(block.get("id"), str) and isinstance(block.get("name"), str) - and isinstance(block.get("input"), Mapping) + and isinstance(block.get("input"), dict) ) @@ -229,22 +230,19 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge assert_never(target) +_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object]) + + def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: content: Final = message.get("content", None) block: Final = content[target.content_idx] if isinstance(content, list) else None if not isinstance(block, dict): return try: - rewritten_input: Final = json.loads(shape.arguments) - except json.JSONDecodeError: + rewritten_input: Final = _TOOL_USE_INPUT_ADAPTER.validate_json(shape.arguments) + except ValidationError: verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", - block.get("id"), - ) - return - if not isinstance(rewritten_input, dict): - verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + "Anthropic Messages: guardrail returned arguments that are not a JSON object for tool_use %s; keeping its input", block.get("id"), ) return @@ -1113,11 +1111,11 @@ class AnthropicMessagesHandler(BaseTranslation): messages: Sequence[_WritableMessage], scanned_tool_calls: tuple[ScannedToolCall, ...], pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], - returned_tool_calls: object, + returned_tool_calls: Sequence[object] | None, ) -> None: post_guardrail_tool_calls: Final = _tool_call_shapes( returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) else tuple(item.tool_call for item in scanned_tool_calls) ) for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): From 330ba7cbf91d4db59f3e1b433aa937fcdaddbb2d Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:21:09 +0000 Subject: [PATCH 076/297] fix(ui): show the team alias on the model info page and in its raw JSON Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/ModelInfoEditForm.tsx | 10 ++- .../src/components/model_info_view.test.tsx | 78 +++++++++++++++++++ .../src/components/model_info_view.tsx | 11 ++- 3 files changed, 96 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index d56e65237eb..fb4d90b5ea2 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -271,6 +271,7 @@ const displayCost = (localModelData: any, field: TouchedPricingField): string => interface ModelInfoEditFormProps { localModelData: any; modelData: { model_info: { team_id?: string | null } & Record }; + teamAlias: string | null; accessToken: string | null; isEditing: boolean; isSaving: boolean; @@ -341,6 +342,7 @@ const ChipList: React.FC<{ values: unknown; emptyLabel: string }> = ({ values, e const ModelInfoEditForm: React.FC = ({ localModelData, modelData, + teamAlias, accessToken, isEditing, isSaving, @@ -799,8 +801,12 @@ const ModelInfoEditForm: React.FC = ({
- Team ID - {modelData.model_info.team_id || "Not Set"} + Team + + {teamAlias + ? `${teamAlias} (${modelData.model_info.team_id})` + : modelData.model_info.team_id || "Not Set"} +
diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 3db9418dfb9..f714b8e5c4a 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -42,6 +42,11 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args), })); +const mockUseTeams = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeams: () => mockUseTeams(), +})); + const mockUsePtuCostAttributionEnabled = vi.fn(); vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({ usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(), @@ -102,6 +107,7 @@ describe("ModelInfoView", () => { }); vi.clearAllMocks(); mockUsePtuCostAttributionEnabled.mockReturnValue(false); + mockUseTeams.mockReturnValue({ data: undefined, isLoading: false, error: null }); mockUseModelsInfo.mockReturnValue({ data: { @@ -1305,6 +1311,78 @@ describe("ModelInfoView", () => { }); }); + describe("team alias", () => { + const teamModel = { + ...defaultModelData, + model_info: { ...defaultModelData.model_info, team_id: "team-1" }, + }; + + beforeEach(() => { + mockUseModelsInfo.mockReturnValue({ data: { data: [teamModel] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [teamModel] }); + }); + + const readRawJson = async (user: ReturnType) => { + await user.click(await screen.findByRole("tab", { name: /raw json/i })); + const pre = await screen.findByText(/"model_name": "GPT-4"/, { selector: "pre" }); + return JSON.parse(pre.textContent ?? ""); + }; + + it("shows the team alias next to the team id and adds team_alias to the raw JSON", async () => { + mockUseTeams.mockReturnValue({ + data: [ + { team_id: "team-0", team_alias: "other" }, + { team_id: "team-1", team_alias: "alpha" }, + ], + isLoading: false, + error: null, + }); + const user = userEvent.setup(); + render(, { wrapper }); + + expect(await screen.findByText("alpha (team-1)")).toBeInTheDocument(); + + const raw = await readRawJson(user); + expect(raw.model_info).toMatchObject({ team_id: "team-1", team_alias: "alpha" }); + const keys = Object.keys(raw.model_info); + expect(keys.indexOf("team_alias")).toBe(keys.indexOf("team_id") + 1); + }); + + it("falls back to the bare team id when the team is not in the caller's team list", async () => { + mockUseTeams.mockReturnValue({ + data: [{ team_id: "team-0", team_alias: "other" }], + isLoading: false, + error: null, + }); + const user = userEvent.setup(); + render(, { wrapper }); + + expect(await screen.findByText("team-1")).toBeInTheDocument(); + + const raw = await readRawJson(user); + expect(raw.model_info.team_id).toBe("team-1"); + expect(raw.model_info).not.toHaveProperty("team_alias"); + }); + + it("shows Not Set and no team_alias for a model without a team", async () => { + mockUseModelsInfo.mockReturnValue({ data: { data: [defaultModelData] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [defaultModelData] }); + mockUseTeams.mockReturnValue({ + data: [{ team_id: "team-1", team_alias: "alpha" }], + isLoading: false, + error: null, + }); + const user = userEvent.setup(); + render(, { wrapper }); + + expect(await screen.findByText("Team")).toBeInTheDocument(); + expect(screen.queryByText(/alpha/)).not.toBeInTheDocument(); + + const raw = await readRawJson(user); + expect(raw.model_info).not.toHaveProperty("team_alias"); + }); + }); + it("renders the provider card logo from the bundled provider map", async () => { render(, { wrapper }); diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 35afcdb2985..f25416327c0 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -169,6 +169,12 @@ export default function ModelInfoView({ // Keep modelData variable name for backwards compatibility const modelData = transformedModelData; + const teamAlias = teams?.find((team) => team.team_id === modelData?.model_info?.team_id)?.team_alias || null; + const rawModelInfoEntries = Object.entries(modelData?.model_info ?? {}).flatMap((entry) => + entry[0] === "team_id" && teamAlias ? [entry, ["team_alias", teamAlias]] : [entry], + ); + const rawModelData = modelData && { ...modelData, model_info: Object.fromEntries(rawModelInfoEntries) }; + const canEditModel = canModifyModel({ userRole, userID, isViewOnly }, teams ?? null, { teamId: modelData?.model_info?.team_id, isDbModel: modelData?.model_info?.db_model === true, @@ -765,6 +771,7 @@ export default function ModelInfoView({ - 
{JSON.stringify(modelData, null, 2)}
+ 
+                {JSON.stringify(rawModelData, null, 2)}
+
From 825e4f17e949439d37f922bb02c0356f2bdd0dc5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:23:34 -0700 Subject: [PATCH 077/297] fix(sdk): carry body and proxy headers on relayed litellm errors and content policy blocks too --- litellm/exceptions.py | 2 + .../exception_mapping_utils.py | 49 +++++++------- .../test_exception_mapping_utils.py | 65 +++++++++++++------ .../common_utils/test_openai_error_payload.py | 2 +- 4 files changed, 72 insertions(+), 46 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3d9e26e450b..fdc2cc1f169 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -626,6 +626,7 @@ class ContentPolicyViolationError(BadRequestError): litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, body: dict | None = None, + headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" @@ -640,6 +641,7 @@ class ContentPolicyViolationError(BadRequestError): response=response, litellm_debug_info=self.litellm_debug_info, body=body, + headers=headers, ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index fd1ba666887..36b53a26c99 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,3 +1,4 @@ +import inspect import json import re import traceback @@ -203,11 +204,18 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None return _response_headers +def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]: + accepted: Final = inspect.signature(exception_class).parameters + return {name: value for name, value in candidates.items() if name in accepted} + + def extract_and_raise_litellm_exception( response: Any | None, error_str: str, model: str, custom_llm_provider: str, + body: object | None = None, + headers: Mapping[str, str] | None = None, ): """ Covers scenario where litellm sdk calling proxy. @@ -217,32 +225,19 @@ def extract_and_raise_litellm_exception( Relevant Issue: https://github.com/BerriAI/litellm/issues/7259 """ pattern: Final = r"litellm\.\w+Error" - - # Search for the exception in the error string match: Final = re.search(pattern, error_str) - - # Extract the exception if found - if match: - exception_name = match.group(0) - exception_name = exception_name.strip().replace("litellm.", "") - raised_exception_obj: Final = getattr(litellm, exception_name, None) - if raised_exception_obj: - # Try with response parameter first, fall back to without it - # Some exceptions (e.g., APIConnectionError) don't accept response param - try: - raise raised_exception_obj( - message=error_str, - llm_provider=custom_llm_provider, - model=model, - response=response, - ) - except TypeError: - # Exception doesn't accept response parameter - raise raised_exception_obj( - message=error_str, - llm_provider=custom_llm_provider, - model=model, - ) + if match is None: + return + exception_name: Final = match.group(0).removeprefix("litellm.") + raised_exception_obj: Final = getattr(litellm, exception_name, None) + if not raised_exception_obj: + return + raise raised_exception_obj( + message=error_str, + llm_provider=custom_llm_provider, + model=model, + **_accepted_init_kwargs(raised_exception_obj, {"response": response, "body": body, "headers": headers}), + ) class _ProviderHTTPException(Protocol): @@ -339,6 +334,8 @@ def _map_openai_exception( model=model, response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: helpful_message: Final = ( @@ -2443,6 +2440,8 @@ def exception_type( error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, + body=getattr(original_exception, "body", None), + headers=_litellm_proxy_response_headers(mappable_exception, custom_llm_provider), ) if ( custom_llm_provider == "openai" diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index ea6ac17ad45..1ff2bbb9bdd 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1424,35 +1424,37 @@ _GUARDRAIL_BLOCK_ERROR = { } -def _openai_handler_error(error_type: str, headers: dict[str, str]) -> OpenAIError: - """What litellm/llms/openai/openai.py raises after the openai SDK rejects a 400: +def _openai_handler_error( + error_type: str, + headers: dict[str, str], + status_code: int = 400, + message: str = _GUARDRAIL_BLOCK_ERROR["message"], +) -> OpenAIError: + """What litellm/llms/openai/openai.py raises after the openai SDK rejects a request: the SDK's str() carries the wire body, and the handler copies headers and body over.""" - wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type} - wire = httpx.Response( - status_code=400, - headers=headers, - json={"error": wire_error}, - request=httpx.Request("POST", "http://localhost:4000/v1/chat/completions"), - ) + wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message} return OpenAIError( - status_code=400, - message=f"Error code: 400 - {{'error': {wire_error}}}", - headers=wire.headers, + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {wire_error}}}", + headers=httpx.Headers(headers), body=wire_error, ) -@pytest.mark.parametrize("error_type", ["None", "invalid_request_error"]) -def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): - """An SDK caller behind a proxy tells a guardrail block from any other 400 by the body's - provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError - must carry both whichever error.type the proxy version on the other end emits.""" - proxy_headers = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} +_PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} + +@pytest.mark.parametrize( + ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] +) +def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): + """An SDK caller behind a proxy tells a guardrail block from any other 4xx by the body's + provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError + must carry both whichever error.type and status the proxy version on the other end emits.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="claude-haiku-4-5", - original_exception=_openai_handler_error(error_type, proxy_headers), + original_exception=_openai_handler_error(error_type, _PROXY_HEADERS, status_code=status_code), custom_llm_provider="litellm_proxy", completion_kwargs={}, extra_kwargs={}, @@ -1460,7 +1462,30 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" assert exc_info.value.body["type"] == error_type - assert proxy_headers.items() <= exc_info.value.headers.items() + assert exc_info.value.headers == _PROXY_HEADERS + + +@pytest.mark.parametrize( + "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] +) +def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): + """A proxy relaying a provider's own litellm error names the class in the message, which + re-raises that class on the SDK side before the generic 400 mapping runs; it must carry the + body and the proxy headers the same way the generic mapping now does.""" + message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" + + with pytest.raises(relayed_class) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error("None", _PROXY_HEADERS, message=message), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert type(exc_info.value) is relayed_class + assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" + assert exc_info.value.headers == _PROXY_HEADERS def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index db9fe3a2253..90f1da84a61 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -145,7 +145,7 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" -def test_the_stringified_none_an_older_upstream_proxy_sent_is_treated_as_absent(): +def test_a_stringified_none_type_or_param_is_treated_as_absent(): """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal is the exact bug this module exists to stop.""" From d0a846c8be5ec1ae9036254eeb04575f6b406921 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:27:33 +0000 Subject: [PATCH 078/297] test(router): cover team_model_has_alternatives directly in the mapped router test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_router.py | 34 +++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..7af60e01f7f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -896,6 +896,40 @@ def test_arouter_test_team_model(): assert result is not None +def test_team_model_has_alternatives(): + def team_deployment(deployment_id: str, team_id: str, public_model_name: str): + return { + "model_name": f"model_name_{team_id}_{deployment_id}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": deployment_id, + "team_id": team_id, + "team_public_model_name": public_model_name, + }, + } + + router = litellm.Router( + model_list=[ + team_deployment("team-a-1", "team-a", "shared-model"), + team_deployment("team-a-2", "team-a", "shared-model"), + team_deployment("team-a-solo", "team-a", "solo-model"), + team_deployment("team-b-1", "team-b", "shared-model"), + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"id": "plain-1"}, + }, + ], + ) + + assert router.team_model_has_alternatives("team-a-1") is True + assert router.team_model_has_alternatives("team-a-2") is True + assert router.team_model_has_alternatives("team-a-solo") is False + assert router.team_model_has_alternatives("team-b-1") is False + assert router.team_model_has_alternatives("plain-1") is False + assert router.team_model_has_alternatives("missing-deployment") is False + + def test_arouter_ignore_invalid_deployments(): """ Test that router.ignore_invalid_deployments is set to True From 37447c98f77116a54e3c902a64a64fae618c7996 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:28:43 -0700 Subject: [PATCH 079/297] fix(guardrails): reject per-message texts that cannot land on a string input or a Messages request --- .../chat/guardrail_translation/handler.py | 3 ++ .../base_llm/guardrail_translation/utils.py | 6 +++ .../chat/guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 7 +-- .../test_anthropic_guardrail_handler.py | 52 +++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 35 +++++++++++++ 6 files changed, 102 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..656e9978eff 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -44,6 +44,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( scoped_structured_message_indices, stream_item_field, stream_item_fingerprint, + unappliable_request_rewrite, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -570,6 +571,8 @@ class AnthropicMessagesHandler(BaseTranslation): preserve_system_messages=has_midturn_system_message, ) else: + if guardrailed_texts and len(guardrailed_texts) != len(scanned): + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( messages=messages, diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 34d648cf184..a80e20c5404 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -411,3 +411,9 @@ def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped + + +def unappliable_request_rewrite(guardrail_name: str | None) -> Exception: + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + return UnappliableRequestRewrite(guardrail_name or "unknown") diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 56fda636e9a..ee68f8f6546 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_field, stream_item_fingerprint, stream_item_items, + unappliable_request_rewrite, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -197,9 +198,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: if len(guardrailed_texts) != len(text_task_mappings): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - - raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) await self._apply_guardrail_responses_to_input_texts( messages=messages, responses=guardrailed_texts, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..27ff55f120c 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_field, stream_item_fingerprint, stream_item_items, + unappliable_request_rewrite, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation): data["instructions"] = written_back.instructions # rebind-ok: data is an out-param elif isinstance(input_data, str): guardrailed_texts: Final = guardrailed_inputs.get("texts") or () + if len(guardrailed_texts) > 1: + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param else: rewritten_texts: Final = guardrailed_inputs.get("texts") or () if len(rewritten_texts) != len(extracted.task_mappings): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - - raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) await self._apply_guardrail_responses_to_input( messages=input_data, responses=rewritten_texts, 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 9fe56f4dc65..d6fd30638cc 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 @@ -2272,6 +2272,58 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert ended_key != open_key +class PerRowTextGuardrail(CustomGuardrail): + """Answers one redacted text per chat row it was shown, the way a guardrail + that scans per message does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="per-row-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + rows = inputs.get("structured_messages") or [] + return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "") for row in rows]} + + +class TestPerMessageTextWriteBack: + """Texts that no longer pair one-to-one with what the handler extracted must be + rejected by name instead of sliding onto the wrong messages.""" + + @pytest.mark.asyncio + async def test_one_text_per_row_over_a_system_prompt_is_rejected_by_name(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + data = { + "model": "claude-sonnet-4-5", + "system": "Reply with exactly the SSN you were given.", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + original = json.loads(json.dumps(data)) + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail()) + + assert excinfo.value.guardrail_name == "per-row-redactor" + assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" + assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" + + @pytest.mark.asyncio + async def test_one_text_per_row_without_a_system_prompt_is_applied(self): + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail()) + + assert data["messages"] == [{"role": "user", "content": "My SSN is ."}] + + class TestAnthropicMessagesHandlerPostCallHookResponse: def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): from litellm.types.utils import Choices, Message, ModelResponse, Usage 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 394f134e99c..ac719da169c 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 @@ -2392,6 +2392,14 @@ def _tool_replay_request() -> dict: } +def _string_input_request() -> dict: + return { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": "My SSN is " + SSN + ".", + } + + class TestPerMessageRewriteWriteBack: """A guardrail that rewrites per chat row hands the rows back as structured_messages, and the handler lands them on the instructions and the @@ -2432,6 +2440,33 @@ class TestPerMessageRewriteWriteBack: assert data["input"] == original["input"] assert data["instructions"] == original["instructions"] + @pytest.mark.asyncio + async def test_structured_rows_land_on_instructions_and_string_input(self): + guardrail = _per_message_redactor() + data = _string_input_request() + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." + assert [_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] + + @pytest.mark.asyncio + async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + guardrail = _per_message_redactor() + data = _string_input_request() + original = copy.deepcopy(data) + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)): + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert excinfo.value.guardrail_name == "per-message-redactor" + assert data["input"] == original["input"] + assert data["instructions"] == original["instructions"] + class TestProvenancePatching: """The O(n) provenance pass must keep patching rewritten rows in place for the From 10f411e60dd7a771a2b3199e3306d197a3127bea Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:34:16 +0000 Subject: [PATCH 080/297] fix(router): name the all-deployments-in-cooldown error on 429 responses RouterRateLimitError now carries the model group's deployment ids so it can tell when every deployment is cooled down, and exposes that as type=all_deployments_in_cooldown with an explicit message. A partial cooldown keeps type=rate_limit_error. Either way the proxy no longer reports type=internal_server_error next to code 429 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 4 ++ litellm/router_utils/handle_error.py | 1 + litellm/types/router.py | 21 ++++++++- .../proxy/test_common_request_processing.py | 33 +++++++++++++ tests/test_litellm/test_router.py | 46 +++++++++++++++++++ 5 files changed, 104 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..0983c9689c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13878,6 +13878,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) if strategy == "simple-shuffle": @@ -13910,6 +13911,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) self._override_selector_pre_call_check(strategy, strategy_selector, deployment) verbose_router_logger.info( @@ -14024,6 +14026,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) # 6. Apply load balancing strategy @@ -14057,6 +14060,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) self._override_selector_pre_call_check(strategy, strategy_selector, deployment) diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 0e7490d31b1..bfe02675162 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -93,4 +93,5 @@ async def async_raise_no_deployment_exception( cooldown_time=_cooldown_time, enable_pre_call_checks=litellm_router_instance.enable_pre_call_checks, cooldown_list=cooldown_list_ids, + model_ids=model_ids, ) diff --git a/litellm/types/router.py b/litellm/types/router.py index c7363502017..ddc13e18567 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -645,6 +645,7 @@ class RouterErrors(enum.Enum): user_defined_ratelimit_error = "Deployment over user-defined ratelimit." no_deployments_available = "No deployments available for selected model" + all_deployments_in_cooldown = "All deployments for selected model are in cooldown" no_deployments_with_tag_routing = "Not allowed to access model due to tags configuration" no_deployments_with_provider_budget_routing = "No deployments available - crossed budget" no_healthy_deployments = "There are no healthy deployments for this model" @@ -868,6 +869,11 @@ class RouterRateLimitErrorBasic(ValueError): super().__init__(_message) +class RouterErrorTypes(str, enum.Enum): + rate_limit_error = "rate_limit_error" + all_deployments_in_cooldown = "all_deployments_in_cooldown" + + class RouterRateLimitError(ValueError): def __init__( self, @@ -875,12 +881,25 @@ class RouterRateLimitError(ValueError): cooldown_time: float, enable_pre_call_checks: bool, cooldown_list: list, + model_ids: Sequence[str] = (), ) -> None: self.model = model self.cooldown_time = cooldown_time self.enable_pre_call_checks = enable_pre_call_checks self.cooldown_list = cooldown_list - _message = f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds. Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}" + self.all_deployments_in_cooldown = bool(model_ids) and frozenset(model_ids) <= frozenset(cooldown_list) + self.type = ( + RouterErrorTypes.all_deployments_in_cooldown.value + if self.all_deployments_in_cooldown + else RouterErrorTypes.rate_limit_error.value + ) + _reason: Final = ( + f" {RouterErrors.all_deployments_in_cooldown.value}." if self.all_deployments_in_cooldown else "" + ) + _message: Final = ( + f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds.{_reason} " + f"Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}" + ) super().__init__(_message) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 69e89d1c604..92b80685a1f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3879,6 +3879,39 @@ class TestHandleLLMApiExceptionRetryAfter: assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" + async def test_handle_llm_api_exception_names_cooldown_when_every_deployment_is_cooled_down(self): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=120, + enable_pre_call_checks=False, + cooldown_list=["dep-a", "dep-b"], + model_ids=["dep-a", "dep-b"], + ) + proxy_exc = await self._invoke(exc) + body = proxy_exc.to_dict() + assert body["type"] == "all_deployments_in_cooldown" + assert body["code"] == "429" + assert "All deployments for selected model are in cooldown" in body["message"] + assert proxy_exc.headers["retry-after"] == "120" + + async def test_handle_llm_api_exception_keeps_rate_limit_type_when_cooldown_is_partial(self): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=120, + enable_pre_call_checks=False, + cooldown_list=["dep-a"], + model_ids=["dep-a", "dep-b"], + ) + proxy_exc = await self._invoke(exc) + body = proxy_exc.to_dict() + assert body["type"] == "rate_limit_error" + assert body["code"] == "429" + assert "All deployments for selected model are in cooldown" not in body["message"] + class TestHandleLLMApiExceptionFramingHeaders: """HTTP-framing headers on the provider exception must be stripped before the diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..cc571ad3d3c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7735,6 +7735,52 @@ def test_get_available_deployment_raises_when_addressed_dict_is_blocked(): router.get_available_deployment(model="dep-0", request_kwargs={}) +def _cool_down(router: Router, *deployment_ids: str) -> None: + for deployment_id in deployment_ids: + router.cooldown_cache.add_deployment_to_cooldown( + model_id=deployment_id, + original_exception=litellm.RateLimitError(message="upstream 429", llm_provider="openai", model="gpt-4o"), + exception_status=429, + cooldown_time=60, + ) + + +async def _select_deployment(router: Router, use_async: bool) -> None: + if use_async: + await router.async_get_available_deployment(model="gpt-4o", request_kwargs={}) + return + router.get_available_deployment(model="gpt-4o", request_kwargs={}) + + +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_get_available_deployment_names_cooldown_when_every_deployment_is_cooled_down(use_async: bool): + from litellm.types.router import RouterErrors, RouterRateLimitError + + router: Final = _router_with_two_deployments([False, False]) + _cool_down(router, "dep-0", "dep-1") + with pytest.raises(RouterRateLimitError) as exc_info: + await _select_deployment(router, use_async) + assert exc_info.value.all_deployments_in_cooldown is True + assert exc_info.value.type == "all_deployments_in_cooldown" + assert RouterErrors.all_deployments_in_cooldown.value in str(exc_info.value) + assert str(exc_info.value).startswith("No deployments available for selected model, Try again in ") + + +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_get_available_deployment_keeps_generic_error_when_cooldown_is_partial(use_async: bool): + from litellm.types.router import RouterErrors, RouterRateLimitError + + router: Final = _router_with_two_deployments([False, True]) + _cool_down(router, "dep-0") + with pytest.raises(RouterRateLimitError) as exc_info: + await _select_deployment(router, use_async) + assert exc_info.value.all_deployments_in_cooldown is False + assert exc_info.value.type == "rate_limit_error" + assert RouterErrors.all_deployments_in_cooldown.value not in str(exc_info.value) + + def _router_with_two_pass_through_deployments(blocked_flags): import litellm From c2463728593b935448bcb09f54ef4fd070e9f8bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:35:53 -0700 Subject: [PATCH 081/297] fix(responses): import BaseLLMException lazily and collect stream chunks via anext Move the BaseLLMException import into _map_error_event_exception so the module no longer imports it at load time, clearing the module-level cyclic import CodeQL flagged. The class is used only on the cold error path. Replace the mutable list-append test collector with aiter/anext so the regression tests read the stream immutably. --- litellm/responses/streaming_iterator.py | 3 ++- .../test_streaming_iterator_error_events.py | 22 ++++++++----------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index b3426fbbfef..a15571acb7c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,7 +31,6 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -553,6 +552,8 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: + from litellm.llms.base_llm.chat.transformation import BaseLLMException + error_message, error_type, error_code = _error_event_fields(error_obj) status_code: Final = _status_code_for_error_fields(error_type, error_code) error_body: Final = {"message": error_message, "type": error_type, "code": error_code} diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 73afbb5e63a..2f4fba45cee 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -287,15 +287,12 @@ async def test_async_iterator_content_policy_violation_after_first_chunk_carries ] ) - chunks = [] - - async def _drain(): - async for chunk in iterator: - chunks.append(chunk) + stream = aiter(iterator) + first_chunk = await anext(stream) + assert first_chunk is not None with pytest.raises(MidStreamFallbackError) as exc_info: - await _drain() - assert len(chunks) == 1 + await anext(stream) assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) assert exc_info.value.is_pre_first_chunk is False assert exc_info.value.generated_content == "partial " @@ -316,14 +313,13 @@ async def test_async_iterator_error_after_first_chunk_carries_generated_content( ] ) - chunks = [] - async def _drain(): - async for chunk in iterator: - chunks.append(chunk) + stream = aiter(iterator) + first_chunk = await anext(stream) + second_chunk = await anext(stream) + assert first_chunk is not None and second_chunk is not None with pytest.raises(MidStreamFallbackError) as exc_info: - await _drain() - assert len(chunks) == 2 + await anext(stream) assert exc_info.value.status_code == 500 assert exc_info.value.is_pre_first_chunk is False assert exc_info.value.generated_content == "hello world" From 9080f0904ad6bee6c5f10debdccdd1b064e1efe3 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:35:54 +0000 Subject: [PATCH 082/297] fix(router): ignore blocked siblings when checking team model cooldown alternatives Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 4 +++- .../router_utils/test_cooldown_handlers.py | 18 +++++++++++++++++- tests/test_litellm/test_router.py | 6 +++++- 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 4929b17f7fc..63d5f86d46a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1561,7 +1561,9 @@ class Router: public_model_name: Final = deployment.model_info.team_public_model_name if team_id is None or public_model_name is None: return False - return len(self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()) > 1 + sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or () + routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices]) + return len(routable_siblings) > 1 _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 6f66bb863cc..fdbfd618dab 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -440,7 +440,7 @@ class TestRoutingGroupCooldownAlternatives: class TestTeamModelCooldownAlternatives: - def _router(self, team_deployments: int): + def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()): from litellm import Router return Router( @@ -452,6 +452,7 @@ class TestTeamModelCooldownAlternatives: "id": f"team-deploy-{i}", "team_id": "team-1", "team_public_model_name": "team-gpt-4o-mini", + "blocked": f"team-deploy-{i}" in blocked_ids, }, } for i in range(team_deployments) @@ -487,3 +488,18 @@ class TestTeamModelCooldownAlternatives: ) is False ) + + def test_429_with_only_a_blocked_sibling_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2, blocked_ids=frozenset({"team-deploy-1"})) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7af60e01f7f..7786bfb4788 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -897,7 +897,7 @@ def test_arouter_test_team_model(): def test_team_model_has_alternatives(): - def team_deployment(deployment_id: str, team_id: str, public_model_name: str): + def team_deployment(deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False): return { "model_name": f"model_name_{team_id}_{deployment_id}", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, @@ -905,6 +905,7 @@ def test_team_model_has_alternatives(): "id": deployment_id, "team_id": team_id, "team_public_model_name": public_model_name, + "blocked": blocked, }, } @@ -914,6 +915,8 @@ def test_team_model_has_alternatives(): team_deployment("team-a-2", "team-a", "shared-model"), team_deployment("team-a-solo", "team-a", "solo-model"), team_deployment("team-b-1", "team-b", "shared-model"), + team_deployment("team-c-1", "team-c", "paused-sibling-model"), + team_deployment("team-c-paused", "team-c", "paused-sibling-model", blocked=True), { "model_name": "plain-model", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, @@ -926,6 +929,7 @@ def test_team_model_has_alternatives(): assert router.team_model_has_alternatives("team-a-2") is True assert router.team_model_has_alternatives("team-a-solo") is False assert router.team_model_has_alternatives("team-b-1") is False + assert router.team_model_has_alternatives("team-c-1") is False assert router.team_model_has_alternatives("plain-1") is False assert router.team_model_has_alternatives("missing-deployment") is False From dd173a0b1b7099255c51812edb327dffddb121fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:35 -0700 Subject: [PATCH 083/297] fix(guardrails): count the PANW latest-user scan over the hoisted system prompt --- .../panw_prisma_airs/panw_prisma_airs.py | 11 +--- .../guardrail_hooks/test_panw_prisma_airs.py | 55 ++++++------------- 2 files changed, 20 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 3bc0dfabefc..9002e2aea07 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): Args: texts: Flattened text entries from the framework. - messages: Original request messages (request_data["messages"]), - NOT structured_messages (which may have injected system content). + messages: The structured messages the framework flattened into ``texts``, + hoisted top-level system prompt included, so positions line up. Returns a set of scannable indices, or None on count mismatch or no user/developer message (safety fallback to existing role-filter behavior). @@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages") if structured_messages: # For Anthropic /v1/messages: default to latest-user-only scanning. - # Uses request_data["messages"] (original format), NOT structured_messages - # (which has injected system content from adapter translation). if self._use_latest_user_only(request_data, logging_obj): - original_messages: Final = request_data.get("messages") - if original_messages: - scannable_indices = self._get_latest_user_text_indices(texts, original_messages) + scannable_indices = self._get_latest_user_text_indices(texts, structured_messages) # Fall through to existing role filtering if: # - not Anthropic, OR flag explicitly False, OR - # - no original messages, OR # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: scannable_indices = self._get_scannable_text_indices(texts, structured_messages) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 3d7c6e06d94..f25727ebd9a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -4620,46 +4620,27 @@ class TestPanwAirsLatestRoleMessageOnly: @pytest.mark.asyncio async def test_anthropic_system_plus_multiturn_no_fallback(self): - """Anthropic with top-level system + multi-turn messages[] - — latest-user works, no scan-all fallback. + """Anthropic with a top-level system prompt and multi-turn messages[] + scans only the latest user turn, with no scan-all fallback. - Key scenario: Anthropic top-level `system` field causes - structured_messages to have an injected system entry, but - request_data["messages"] does NOT include it. + The Anthropic handler hoists the top-level `system` field into both + `texts` and `structured_messages`, so the latest-user walk has to + count the same entries the framework flattened. """ - handler = PanwPrismaAirsHandler( - guardrail_name="test_panw_airs", - api_key="test_api_key", - profile_name="test_profile", - default_on=True, + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, ) - # Original Anthropic messages (no system in messages array) - original_messages = [ - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - # texts extracted from original_messages (3 text entries) - texts = ["First user turn", "First assistant turn", "Latest user turn"] - - # structured_messages has an INJECTED system message from translation - structured_messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - inputs: GenericGuardrailAPIInputs = { - "texts": texts, - "structured_messages": structured_messages, - } + handler = make_handler() request_data = { "litellm_call_id": "test-call-id", "model": "anthropic/claude-sonnet-4-20250514", - "messages": original_messages, + "system": "You are a helpful assistant.", + "messages": [ + {"role": "user", "content": "First user turn"}, + {"role": "assistant", "content": "First assistant turn"}, + {"role": "user", "content": "Latest user turn"}, + ], "proxy_server_request": { "url": "http://localhost:4000/v1/messages", }, @@ -4670,13 +4651,11 @@ class TestPanwAirsLatestRoleMessageOnly: ) as mock_api: mock_api.return_value = {"action": "allow", "category": "benign"} - await handler.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", + await AnthropicMessagesHandler().process_input_messages( + data=request_data, + guardrail_to_apply=handler, ) - # Should scan ONLY the latest user message, not fall back to scan-all assert mock_api.call_count == 1 assert mock_api.call_args.kwargs["content"] == "Latest user turn" From 6264bd84bf1e0b33865028acdf53f5d64adb1e98 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:36 -0700 Subject: [PATCH 084/297] chore(ui): regenerate dashboard API types --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From e61b6bfd5ff58c61381e32aae075b2516fd3c76d Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:48:35 +0000 Subject: [PATCH 085/297] fix(router): classify pass-through cooldown against pass-through deployments only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 7 ++++++- tests/test_litellm/test_router.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 0983c9689c3..e74b2079fd4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13989,6 +13989,11 @@ class Router: model=model, llm_provider="", ) + pass_through_model_ids: Final = tuple( + deployment["model_info"]["id"] + for deployment in pass_through_deployments + if "id" in deployment.get("model_info", {}) + ) # 4. Apply health-check and cooldown filtering parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs) @@ -14026,7 +14031,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, - model_ids=model_ids, + model_ids=pass_through_model_ids, ) # 6. Apply load balancing strategy diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cc571ad3d3c..35d0b8104bd 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7818,6 +7818,24 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): ) +def test_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): + from litellm.types.router import RouterRateLimitError + + router: Final = _router_with_two_pass_through_deployments([False, False]) + router.add_deployment( + Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-plain", api_key="sk-fake-for-tests"), + model_info=ModelInfo(id="plain-0"), + ) + ) + _cool_down(router, "pt-0", "pt-1") + with pytest.raises(RouterRateLimitError) as exc_info: + router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) + assert exc_info.value.all_deployments_in_cooldown is True + assert exc_info.value.type == "all_deployments_in_cooldown" + + def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): """ Bedrock deployments using IAM/OIDC auth have no api_key; pass-through From e41b3bd13fa6415de7b4076dda82f673fac8b957 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:50:12 +0000 Subject: [PATCH 086/297] test(router): annotate return types of team cooldown test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/router_utils/test_cooldown_handlers.py | 6 ++---- tests/test_litellm/test_router.py | 4 +++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index fdbfd618dab..6fed4be5909 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -440,10 +440,8 @@ class TestRoutingGroupCooldownAlternatives: class TestTeamModelCooldownAlternatives: - def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()): - from litellm import Router - - return Router( + def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()) -> litellm.Router: + return litellm.Router( model_list=[ { "model_name": f"model_name_team-1_{i}", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7786bfb4788..c6b3b7fb8d6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -897,7 +897,9 @@ def test_arouter_test_team_model(): def test_team_model_has_alternatives(): - def team_deployment(deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False): + def team_deployment( + deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False + ) -> DeploymentTypedDict: return { "model_name": f"model_name_{team_id}_{deployment_id}", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, From db79226b6b8786b40d10e8736595a0fe6bf07f47 Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 13 Sep 2026 09:59:39 +0000 Subject: [PATCH 087/297] test(auth): freeze the cache clock in auth prefetch tests The org cache entries written by prefetch_auth_objects carry the 5s DEFAULT_IN_MEMORY_TTL. The first @log_db_metrics getter lazily imports litellm.proxy.proxy_server, which on a cold CI runner can take longer than 5s, so the org entry expired before get_org_object read it and the getter fell through to the MagicMock database. Inject a frozen clock into InMemoryCache so the test asserts the join, not import latency. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy_behavior/auth/test_auth_object_prefetch.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py index e2d947f4284..59e9585a296 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -22,6 +22,11 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache pytestmark = pytest.mark.asyncio(loop_scope="session") +def _frozen_cache() -> UserApiKeyCache: + """The org entries carry a 5s TTL; a frozen clock keeps a slow first call from expiring them mid-test.""" + return UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1_000_000.0), redis_cache=None) + + def _dead_db() -> MagicMock: prisma = MagicMock(name="prisma_client") prisma.db.query_first = AsyncMock(return_value=None) @@ -58,7 +63,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma): data={"user_id": user_id, "team_id": team_b, "litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}}} ) - cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) @@ -100,7 +105,7 @@ async def test_join_reads_team_model_aliases_from_the_mapped_column(prisma): where={"team_id": team_id}, include={"litellm_model_table": True} ) - cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + cache = _frozen_cache() refs = AuthObjectRefs(user_id=None, team_id=team_id, membership_user_id=None, organization_id=None) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) @@ -144,7 +149,7 @@ async def test_join_reads_null_nested_lists_the_way_prisma_does(prisma): where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True} ) - cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_id, membership_user_id=user_id, organization_id=None) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) From 7ea19eccc7a6ebe4f6c0bab35073f02b17407fe5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:02:03 -0700 Subject: [PATCH 088/297] fix(responses): keep namespace custom tools through guardrail merges and Mantle params identity --- .../responses/transformation.py | 9 +++-- .../guardrail_translation/tool_merge.py | 31 ++++++++-------- .../custom_tools.py | 35 +++++++++++++------ .../transformation.py | 3 +- ...bedrock_mantle_responses_transformation.py | 6 ++++ ...t_openai_responses_guardrail_tool_merge.py | 32 +++++++++++++++-- 6 files changed, 85 insertions(+), 31 deletions(-) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 95375399033..57590601a3c 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -237,10 +237,15 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) hoisted: Final = hoist_additional_tools(input, params.get("tools")) normalized_input: Final = self._normalize_codex_input_items(hoisted.input) + request_params: Final = ( + self._params_with_hoisted_tools(params, hoisted) + if hoisted.hoisted + else response_api_optional_request_params + ) return super().transform_responses_api_request( model=model, input=normalized_input, - response_api_optional_request_params=self._params_with_hoisted_tools(params, hoisted), + response_api_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -249,8 +254,6 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def _params_with_hoisted_tools( cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools ) -> dict[str, object]: - if not hoisted.hoisted: - return dict(params) supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools)) if supported_tools: return {**params, "tools": supported_tools} diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index b596adfad6f..0326e9b2bfd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -6,8 +6,10 @@ from typing import Final, TypeAlias from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix from litellm.responses.litellm_completion_transformation.transformation import ( NAMESPACE_DESCRIPTION_SEPARATOR, + NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS, LiteLLMCompletionResponsesConfig, ) @@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]: return tuple(tool for tool in validated if tool is not None) -def _is_function(tool: Tool) -> bool: - return tool.get("type") == "function" +def _has_chat_tool(member: Tool) -> bool: + return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS def _chat_tool_key(tool: Tool) -> str: @@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool: return function if function is not None else MappingProxyType({}) -def _without_namespace_prefix(key: str, value: object, prefix: str) -> object: - if key != "description" or not isinstance(value, str) or not value.startswith(prefix): +def _member_description(key: str, value: object, prefix: str, suffix: str) -> object: + if key != "description" or not isinstance(value, str): return value - return value[len(prefix) :] + return value.removeprefix(prefix).removesuffix(suffix) def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool: flattened_function: Final = _function_fields(flattened) prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else "" + suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else "" changed_function: Final = MappingProxyType( { - key: _without_namespace_prefix(key, value, prefix) + key: _member_description(key, value, prefix, suffix) for key, value in _function_fields(guardrailed).items() if flattened_function.get(key) != value } @@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_ return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType -def _rebuilt_function_members( - function_members: Sequence[Tool], +def _rebuilt_flattened_members( + flattened_members: Sequence[Tool], flattened_group: Sequence[Tool], group_keys: Sequence[IndexedKey], guardrailed_by_key: Mapping[IndexedKey, Tool], @@ -106,7 +109,7 @@ def _rebuilt_function_members( else member if guardrailed_by_key[key] == flattened else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description) - for member, flattened, key in zip(function_members, flattened_group, group_keys) + for member, flattened, key in zip(flattened_members, flattened_group, group_keys) ) @@ -118,9 +121,9 @@ def _rebuilt_namespace( guardrailed_by_key: Mapping[IndexedKey, Tool], ) -> tuple[Tool, ...]: namespace_description: Final = str(original.get("description") or "") - rebuilt_functions: Final = iter( - _rebuilt_function_members( - tuple(member for member in members if _is_function(member)), + rebuilt_flattened: Final = iter( + _rebuilt_flattened_members( + tuple(member for member in members if _has_chat_tool(member)), flattened_group, group_keys, guardrailed_by_key, @@ -129,7 +132,7 @@ def _rebuilt_namespace( ) rebuilt_members: Final = tuple( rebuilt - for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members) + for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members) if rebuilt is not None ) if not rebuilt_members: @@ -149,7 +152,7 @@ def _merged_original( if guardrailed_group == tuple(flattened_group): return (original,) members: Final = _namespace_members(original) if original.get("type") == "namespace" else () - if members and sum(map(_is_function, members)) == len(flattened_group): + if members and sum(map(_has_chat_tool, members)) == len(flattened_group): return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key) if not guardrailed_group: return () diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index 038964055c3..7888a07e248 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -45,21 +45,32 @@ class _ToolNameFields(BaseModel): tools: tuple[object, ...] = () -def _custom_tool_names_of(tool: object) -> tuple[str, ...]: +def _tool_name_fields_of(tool: object) -> _ToolNameFields | None: try: - parsed: Final = _ToolNameFields.model_validate(tool) + return _ToolNameFields.model_validate(tool) except ValidationError: + return None + + +def _custom_tool_name_of(tool: object) -> str | None: + parsed: Final = _tool_name_fields_of(tool) + if parsed is None or parsed.type != "custom" or not parsed.name: + return None + return parsed.name + + +def _nested_tools_of(tool: object) -> tuple[object, ...]: + parsed: Final = _tool_name_fields_of(tool) + if parsed is None or parsed.type != "namespace": return () - if parsed.type == "custom": - return (parsed.name,) if parsed.name else () - if parsed.type != "namespace": - return () - return tuple(name for nested_tool in parsed.tools for name in _custom_tool_names_of(nested_tool)) + return parsed.tools def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]: - """Extract names of tools defined as ``type: "custom"``, at the top level or inside a ``namespace`` tool.""" - return {name for tool in tools or () for name in _custom_tool_names_of(tool)} + """Extract names of ``type: "custom"`` tools, at the top level or one level inside a ``namespace`` tool.""" + top_level: Final = tuple(tools or ()) + nested: Final = tuple(nested_tool for tool in top_level for nested_tool in _nested_tools_of(tool)) + return {name for tool in (*top_level, *nested) if (name := _custom_tool_name_of(tool)) is not None} def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: @@ -155,7 +166,7 @@ def validated_allowed_callers(value: object) -> list[str] | None: raise ValueError("allowed_callers must be a list of strings") from exc -def _grammar_suffix(fmt: object) -> str: +def custom_tool_grammar_suffix(fmt: object) -> str: try: parsed: Final = _CustomToolFormat.model_validate(fmt) except ValidationError: @@ -179,7 +190,9 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp raw_name: Final = tool.get("name") name: Final = raw_name if isinstance(raw_name, str) else "" raw_description: Final = tool.get("description") - description = (raw_description if isinstance(raw_description, str) else "") + _grammar_suffix(tool.get("format")) + description: Final = (raw_description if isinstance(raw_description, str) else "") + custom_tool_grammar_suffix( + tool.get("format") + ) allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers")) function_chunk: Final = ChatCompletionToolParamFunctionChunk( name=name, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 27756b405ec..d27fc855be6 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -109,6 +109,7 @@ NamespaceTool: TypeAlias = Mapping[str, object] ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" +NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) @dataclass(frozen=True, slots=True) @@ -1891,7 +1892,7 @@ class LiteLLMCompletionResponsesConfig: nested: bool, ) -> ChatCompletionToolParam | None: tool_type: Final = namespace_tool.get("type") - if nested and tool_type not in ("function", "custom"): + if nested and tool_type not in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: return None raw_description: Final = str(namespace_tool.get("description") or "") diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 40566261c84..a7aefa714aa 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -773,6 +773,12 @@ class TestBedrockMantleCodexAdditionalTools: assert body["input"] == codex_agentic_items assert "tools" not in body + def test_input_without_additional_tools_sanitizes_tools_on_the_caller_params_object(self): + params = {"tools": [{"type": "function", "name": "wait", "parameters": '{"type": "object"}'}]} + body = self._transform(input=[self._USER_MESSAGE], params=params) + assert body["tools"][0]["parameters"] == {"type": "object"} + assert params["tools"][0]["parameters"] == {"type": "object"} + def test_malformed_additional_tools_item_without_tools_list_is_stripped(self): body = self._transform( input=[ diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index 9c236d81f51..a7f65545eb2 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -133,7 +133,20 @@ def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited( assert merged[0]["tools"][1] == custom_member -def test_namespace_keeps_its_non_function_members_when_every_function_member_is_dropped(): +def test_namespace_keeps_its_custom_member_when_every_function_member_is_dropped(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, + _function("a"), + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [groups[0][1], groups[1][0]]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + + +def test_namespace_custom_member_is_dropped_when_the_guardrail_drops_its_chat_form(): custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} original = [ {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, @@ -143,7 +156,22 @@ def test_namespace_keeps_its_non_function_members_when_every_function_member_is_ merged = merge_guardrailed_tools(original, groups, [groups[1][0]]) - assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + assert list(merged) == [_function("a")] + + +def test_custom_member_description_edit_lands_without_the_namespace_prefix_or_grammar_block(): + grammar = {"type": "grammar", "syntax": "lark", "definition": "start: X"} + custom_member = {"type": "custom", "name": "exec", "description": "Run a command", "format": grammar} + original = [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [custom_member]}] + groups = _groups(original) + assert groups[0][0]["function"]["description"] == "Shell\n\nRun a command\n\nFormat:\n```lark\nstart: X\n```" + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "Shell\n\nRun a command (guarded)\n\nFormat:\n```lark\nstart: X\n```" + + merged = merge_guardrailed_tools(original, groups, edited) + + guarded_member = {**custom_member, "description": "Run a command (guarded)"} + assert list(merged) == [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [guarded_member]}] def test_member_extras_edited_by_the_guardrail_land_on_that_member(): From 2923c4ac5520642bc3c3e728f9b2974850a43f7e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:02:15 -0700 Subject: [PATCH 089/297] fix(guardrails): read the rewrite from texts when a guardrail echoes every row back unchanged --- .../generic_guardrail_api.py | 25 ++++++++------ .../test_generic_guardrail_api.py | 34 +++++++++++++++++++ 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16159d32a7f..3d1a173635e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -150,16 +150,20 @@ def _extract_inbound_headers( return None -def _rows_with_unchanged_originals( +def _structured_rows_to_write_back( original_rows: Sequence[AllMessageValues] | None, shown_rows: Sequence[AllMessageValues] | None, returned_rows: Sequence[AllMessageValues], -) -> tuple[AllMessageValues, ...]: +) -> tuple[AllMessageValues, ...] | None: """The request model drops row keys its message types do not declare, so a - row the server echoes back verbatim is restored to the original row object; - only rows the server actually changed reach the endpoint write-back.""" + row the server echoes back verbatim is restored to the original row object. + A server that echoes every row back unchanged has not rewritten anything + per row, so its answer is read from texts, as it was before rows could be + returned at all.""" if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): return tuple(returned_rows) + if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)): + return None return tuple( original if returned == shown else returned for original, shown, returned in zip(original_rows, shown_rows, returned_rows) @@ -354,12 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools - if guardrail_response.structured_messages: - return_inputs["structured_messages"] = list( # mutable-ok: guardrail inputs take a list - _rows_with_unchanged_originals( - structured_messages, shown_messages, guardrail_response.structured_messages - ) - ) + rows_to_write_back: Final = ( + _structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages) + if guardrail_response.structured_messages + else None + ) + if rows_to_write_back is not None: + return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs 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 cc9942e0e40..a5e79f84ef1 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 @@ -659,6 +659,40 @@ class TestStructuredMessagesInResponse: assert returned_rows[1] is tool_call_row assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + @pytest.mark.asyncio + async def test_rows_all_echoed_back_as_shown_leave_the_rewrite_to_texts( + self, generic_guardrail, mock_request_data_input + ): + """A server written against the texts contract that echoes the request rows back + untouched while rewriting texts still gets its texts rewrite applied.""" + original_rows = [ + {"role": "system", "content": "Never repeat an SSN."}, + {"role": "user", "content": "Look up 123-45-6789 for me."}, + ] + + def echo_rows_and_rewrite_texts(url, json, headers): + answer = MagicMock() + answer.json.return_value = { + "action": "NONE", + "texts": [text.replace("123-45-6789", "[REDACTED]") for text in json["texts"]], + "structured_messages": json["structured_messages"], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_rows_and_rewrite_texts): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={ + "texts": ["Never repeat an SSN.", "Look up 123-45-6789 for me."], + "structured_messages": original_rows, + }, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."] + @pytest.mark.asyncio @pytest.mark.parametrize( "structured_messages", From fff7a2cecfb113082681665d236525f2690aae45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:18:59 -0700 Subject: [PATCH 090/297] fix(responses): keep context-window events out of mid-stream fallback and fix stale exception assertions --- litellm/responses/streaming_iterator.py | 2 +- litellm/router.py | 7 +++---- .../test_openai_responses_api.py | 9 +++++---- ...est_router_aresponses_streaming_fallback.py | 2 +- .../test_streaming_iterator_error_events.py | 18 ++++++++++-------- 5 files changed, 20 insertions(+), 18 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index a15571acb7c..b39e130242d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -222,7 +222,7 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: - if isinstance(mapped_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError)): + if isinstance(mapped_exception, litellm.ContentPolicyViolationError): return True status_code: Final = getattr(mapped_exception, "status_code", None) return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 diff --git a/litellm/router.py b/litellm/router.py index 79854a11150..a9a8f3e2739 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3271,12 +3271,11 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. fallback_trigger: Final[Exception] = ( e.original_exception - if isinstance( - e.original_exception, - (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), - ) + if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e ) fallback_response = await self.async_function_with_fallbacks_common_utils( diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index 05bb9113835..c7712d96969 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1627,9 +1627,10 @@ async def test_openai_responses_api_token_limit_error(): Parsing the in-stream ErrorEvent must not raise "pydantic_core._pydantic_core.ValidationError: 3 validation errors for ErrorEvent". - The iterator now surfaces the event as litellm.APIError with status 400 - (invalid_request_error is a non-retriable client error, so no - MidStreamFallbackError wrapping) carrying the provider's message. + The iterator routes the event through litellm.exception_type, so it surfaces as + the typed 400 client error the non-streaming path raises (litellm.BadRequestError) + carrying the provider's message. invalid_request_error is a non-retriable client + error, so there is no MidStreamFallbackError wrapping. """ litellm._turn_on_debug() @@ -1644,7 +1645,7 @@ async def test_openai_responses_api_token_limit_error(): async for event in response: print(event) - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: await _drain() assert exc_info.value.status_code == 400 diff --git a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py index ee4750e9db8..0d33435cf7a 100644 --- a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py +++ b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py @@ -372,7 +372,7 @@ async def test_aresponses_fallback_on_in_stream_error_event(): raised = mock_fallback.await_args.kwargs["e"] assert isinstance(raised, MidStreamFallbackError) assert raised.status_code == 429 - assert isinstance(raised.original_exception, litellm.APIError) + assert isinstance(raised.original_exception, litellm.RateLimitError) assert raised.original_exception.status_code == 429 assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question" diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 2f4fba45cee..3d7c220804a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -4,11 +4,11 @@ raise instead of being returned as benign chunks, mirroring chat streaming semantics (_handle_stream_fallback_error). The event's code, type and status go through litellm.exception_type, so each event raises the same typed exception the non-streaming path raises for that provider error: non-retriable 4xx -(except 429) raise that typed exception directly, while 429, 5xx, -ContentPolicyViolationError and ContextWindowExceededError are wrapped in +(except 429) raise that typed exception directly, so a context-length event +surfaces as ContextWindowExceededError(400) with no MidStreamFallbackError +wrapping, while 429, 5xx and ContentPolicyViolationError are wrapped in MidStreamFallbackError so the Router's mid-stream fallback machinery fires and -its content_policy_fallbacks / context_window_fallbacks dispatch sees the -trigger it matches on. +its content_policy_fallbacks dispatch sees the trigger it matches on. Status mapping must consider both the OpenAI error `type` (e.g. "invalid_request_error") and `code` (e.g. "invalid_prompt", @@ -110,19 +110,21 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): assert not isinstance(exc_info.value, MidStreamFallbackError) -def test_maybe_raise_for_error_event_wraps_context_window_exceeded_for_context_window_fallbacks(): +def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(): """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming - path and is wrapped so the Router's context_window_fallbacks dispatch fires mid-stream.""" + path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream + fallback, preserving the direct-SDK 400 contract from issue #15785.""" iterator = _make_iterator() chunk = _make_error_chunk( "invalid_request_error", "context_length_exceeded", "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", ) - with pytest.raises(MidStreamFallbackError) as exc_info: + with pytest.raises(litellm.ContextWindowExceededError) as exc_info: iterator._maybe_raise_for_error_event(chunk) - assert isinstance(exc_info.value.original_exception, litellm.ContextWindowExceededError) assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert "maximum context length" in str(exc_info.value) CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." From 0201ca60e7397912cf46e0a1bddf1d6befb86e0a Mon Sep 17 00:00:00 2001 From: ryan Date: Sun, 13 Sep 2026 23:09:27 +0000 Subject: [PATCH 091/297] fix(ui): move tags typed into key metadata JSON into the Tags field Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../templates/KeyEditViewControls.tsx | 41 +++++++++++++++++- .../templates/keyEditFieldNormalizers.test.ts | 42 +++++++++++++++++++ .../templates/keyEditFieldNormalizers.ts | 33 +++++++++++++++ .../key_edit_view.integration.test.tsx | 26 ++++++++++++ .../components/templates/key_edit_view.tsx | 22 ++++++---- 5 files changed, 156 insertions(+), 8 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.test.ts diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index f5d5562ccfe..39542882798 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -1,12 +1,15 @@ import React from "react"; -import { Control } from "react-hook-form"; +import { Control, UseFormReturn } from "react-hook-form"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { CircleHelp } from "lucide-react"; import { FormField } from "@/components/shared/form/FormField"; +import { toast } from "@/lib/toast"; import AgentSelector from "../agent_management/AgentSelector"; import NumericalInput from "../shared/numerical_input"; import SkillSelector from "../skills/SkillSelector"; +import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; import { AgentsAndGroups, KeyEditFormValues } from "./keyEditFormValues"; export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => ( @@ -85,6 +88,42 @@ export const KeyAgentAndSkillFields = ({ ); +type KeyEditForm = Pick< + UseFormReturn, + "control" | "getValues" | "setValue" +>; + +export const moveMetadataTagsToTagsField = (form: KeyEditForm): void => { + const moved = moveTagsOutOfMetadataJson(form.getValues("metadata"), form.getValues("tags")); + if (moved === null) return; + form.setValue("metadata", moved.metadata, { shouldDirty: true }); + form.setValue("tags", moved.tags, { shouldDirty: true }); + if (moved.movedTags.length > 0) { + toast.info(`Moved ${moved.movedTags.join(", ")} from metadata to the Tags field`); + } +}; + +export const KeyMetadataField = ({ form }: { form: KeyEditForm }) => ( + + {(field) => ( +