From e4a047526334dc97a93ef364878b14760e85408d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 25 Jul 2026 10:57:19 -0700 Subject: [PATCH 001/245] test: remove four mirror test files that exercise none of their module A second mutation batch scored the previously unmapped mirror files on current staging. These four generate mutants for the module they are named after, yet no test in the file executes any of them; their test-context coverage lands on generic shared machinery or, for the guardrail translation handler remainder, on no litellm line at all. Eight sibling findings that do exercise a different real module are kept for retargeting instead of removal. --- .../datadog/test_datadog_llm_observability.py | 1195 ----------------- .../guardrail_translation/test_handler.py | 37 - .../test_reasoning_content_transformation.py | 296 ---- tests/test_litellm/test_azure_video_router.py | 53 - 4 files changed, 1581 deletions(-) delete mode 100644 tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py delete mode 100644 tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py delete mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py delete mode 100644 tests/test_litellm/test_azure_video_router.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py deleted file mode 100644 index 1cc3591392b..00000000000 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py +++ /dev/null @@ -1,1195 +0,0 @@ -import asyncio -import os -import sys -from datetime import datetime, timedelta, timezone -from typing import Optional -from unittest.mock import MagicMock, Mock, patch - -import pytest - -# Adds the grandparent directory to sys.path to allow importing project modules -sys.path.insert(0, os.path.abspath("../..")) -import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger -from litellm.types.integrations.datadog_llm_obs import ( - DatadogLLMObsInitParams, -) -from litellm.types.utils import ( - StandardLoggingGuardrailInformation, - StandardLoggingHiddenParams, - StandardLoggingMetadata, - StandardLoggingModelInformation, - StandardLoggingPayload, - StandardLoggingPayloadErrorInformation, -) - - -def create_standard_logging_payload_with_cache() -> StandardLoggingPayload: - """Create a real StandardLoggingPayload object for testing""" - return StandardLoggingPayload( - id="test-request-id-456", - call_type="completion", - response_cost=0.05, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=10, - completion_tokens=20, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4", model_map_value=None - ), - model="gpt-4", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=True, - cache_key="test-cache-key-789", - saved_cache_cost=0.02, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - model_parameters={"stream": True}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key="test-cache-key-789", - api_base="https://api.openai.com", - response_cost="0.05", - additional_headers=None, - ), - trace_id="test-trace-id-123", - custom_llm_provider="openai", - ) - - -def create_standard_logging_payload_with_failure() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object for failure testing""" - return StandardLoggingPayload( - id="test-request-id-failure-789", - call_type="completion", - response_cost=0.0, - response_cost_failure_debug_info=None, - status="failure", - total_tokens=0, - prompt_tokens=10, - completion_tokens=0, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4", model_map_value=None - ), - model="gpt-4", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response=None, - error_str="RateLimitError: You exceeded your current quota", - error_information=StandardLoggingPayloadErrorInformation( - error_code="rate_limit_exceeded", - error_class="RateLimitError", - llm_provider="openai", - traceback="Traceback (most recent call last):\n File test.py, line 1\n RateLimitError: You exceeded your current quota", - error_message="RateLimitError: You exceeded your current quota", - ), - model_parameters={"stream": False}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key=None, - api_base="https://api.openai.com", - response_cost="0.0", - additional_headers=None, - ), - trace_id="test-trace-id-failure-456", - custom_llm_provider="openai", - ) - - -class TestDataDogLLMObsLogger: - """Test suite for DataDog LLM Observability Logger""" - - @pytest.fixture - def mock_env_vars(self): - """Mock environment variables for DataDog""" - with patch.dict( - os.environ, {"DD_API_KEY": "test_api_key", "DD_SITE": "us5.datadoghq.com"} - ): - yield - - @pytest.fixture - def mock_response_obj(self): - """Create a mock response object""" - mock_response = Mock() - mock_response.__getitem__ = Mock( - return_value={ - "choices": [ - { - "message": Mock( - json=Mock( - return_value={"role": "assistant", "content": "Hello!"} - ) - ) - } - ] - } - ) - return mock_response - - def test_cost_and_trace_id_integration(self, mock_env_vars, mock_response_obj): - """Test that total_cost is passed and trace_id from standard payload is used""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_cache() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": { - "metadata": {"trace_id": "old-trace-id-should-be-ignored"} - }, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Test 1: Verify total_cost is correctly extracted from response_cost - assert payload["metrics"].get("total_cost") == 0.05 - - # Test 2: Verify trace_id comes from standard_logging_payload, not metadata - assert payload["trace_id"] == "test-trace-id-123" - - # Test 3: Verify saved_cache_cost is in metadata - metadata = payload["meta"]["metadata"] - assert metadata["saved_cache_cost"] == 0.02 - assert metadata["cache_hit"] is True - assert metadata["cache_key"] == "test-cache-key-789" - - # Test 4: Verify is_streamed_request is in metadata - assert metadata["is_streamed_request"] is True - - def test_cache_metadata_fields(self, mock_env_vars, mock_response_obj): - """Test that cache-related metadata fields are correctly tracked""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_cache() - - # Test the _get_dd_llm_obs_payload_metadata method directly - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - - # Verify all cache-related fields are present - assert metadata["cache_hit"] is True - assert metadata["cache_key"] == "test-cache-key-789" - assert metadata["saved_cache_cost"] == 0.02 - assert metadata["id"] == "test-request-id-456" - assert metadata["trace_id"] == "test-trace-id-123" - assert metadata["model_name"] == "gpt-4" - assert metadata["model_provider"] == "openai" - - def test_get_time_to_first_token_seconds(self, mock_env_vars): - """Test the _get_time_to_first_token_seconds method for streaming calls""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test streaming case (completion_start_time available) - streaming_payload = create_standard_logging_payload_with_cache() - # Modify times for testing: start=1000, completion_start=1002, end=1005 - streaming_payload["startTime"] = 1000.0 - streaming_payload["completionStartTime"] = 1002.0 - streaming_payload["endTime"] = 1005.0 - - # Test streaming case: should use completion_start_time - start_time - time_to_first_token = logger._get_time_to_first_token_seconds( - streaming_payload - ) - assert time_to_first_token == 2.0 # 1002.0 - 1000.0 = 2.0 seconds - - def test_datadog_span_kind_mapping(self, mock_env_vars): - """Test that call_type values are correctly mapped to DataDog span kinds""" - from litellm.types.utils import CallTypes - - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test embedding operations - assert ( - logger._get_datadog_span_kind(CallTypes.embedding.value, "123") - == "embedding" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.aembedding.value, "123") - == "embedding" - ) - - # Test LLM completion operations - assert logger._get_datadog_span_kind(CallTypes.completion.value, None) == "llm" - assert logger._get_datadog_span_kind(CallTypes.acompletion.value, None) == "llm" - assert ( - logger._get_datadog_span_kind(CallTypes.text_completion.value, None) - == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.generate_content.value, None) - == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.anthropic_messages.value, None) - == "llm" - ) - assert logger._get_datadog_span_kind(CallTypes.responses.value, None) == "llm" - assert logger._get_datadog_span_kind(CallTypes.aresponses.value, None) == "llm" - - # Test tool operations - assert ( - logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, "123") - == "tool" - ) - - # Test retrieval operations - assert ( - logger._get_datadog_span_kind(CallTypes.get_assistants.value, "123") - == "retrieval" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.file_retrieve.value, "123") - == "retrieval" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.retrieve_batch.value, "123") - == "retrieval" - ) - - # Test task operations - assert ( - logger._get_datadog_span_kind(CallTypes.create_batch.value, "123") == "task" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.image_generation.value, "123") - == "task" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.moderation.value, "123") == "task" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.transcription.value, "123") - == "task" - ) - - # Test default fallback - assert logger._get_datadog_span_kind("unknown_call_type", None) == "llm" - assert logger._get_datadog_span_kind(None, None) == "llm" - - def test_datadog_span_kind_defaults_without_parent(self, mock_env_vars): - """Test that non-llm kinds fallback to llm when no parent span is provided""" - from litellm.types.utils import CallTypes - - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Tool/task/retrieval span kinds should fallback to llm when parent_id missing - assert ( - logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, None) == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.create_batch.value, None) == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.get_assistants.value, None) == "llm" - ) - - @pytest.mark.asyncio - async def test_async_log_failure_event(self, mock_env_vars): - """Test that async_log_failure_event correctly processes failure payloads according to DD LLM Obs API spec""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Ensure log_queue starts empty - logger.log_queue = [] - - standard_failure_payload = create_standard_logging_payload_with_failure() - - kwargs = { - "standard_logging_object": standard_failure_payload, - "model": "gpt-4", - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() + timedelta(seconds=2) - - # Mock async_send_batch to prevent actual network calls - with patch.object(logger, "async_send_batch") as mock_send_batch: - # Call the method under test - await logger.async_log_failure_event(kwargs, None, start_time, end_time) - - # Verify payload was added to queue - assert len(logger.log_queue) == 1 - - # Verify the payload has correct failure characteristics according to DD LLM Obs API spec - payload = logger.log_queue[0] - assert payload["trace_id"] == "test-trace-id-failure-456" - assert ( - payload["meta"]["metadata"]["id"] == "test-request-id-failure-789" - ) - assert payload["status"] == "error" - - # Verify error information follows DD LLM Obs API spec - assert ( - payload["meta"]["error"]["message"] - == "RateLimitError: You exceeded your current quota" - ) - assert payload["meta"]["error"]["type"] == "RateLimitError" - assert ( - payload["meta"]["error"]["stack"] - == "Traceback (most recent call last):\n File test.py, line 1\n RateLimitError: You exceeded your current quota" - ) - - assert payload["metrics"]["total_cost"] == 0.0 - assert payload["metrics"]["total_tokens"] == 0 - assert payload["metrics"]["output_tokens"] == 0 - - # Verify batch sending not triggered (queue size < batch_size) - mock_send_batch.assert_not_called() - - -class TestDataDogLLMObsLoggerForRedaction(DataDogLLMObsLogger): - """Test suite for DataDog LLM Observability Logger""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.logged_standard_logging_payload = kwargs.get("standard_logging_object") - - -class TestS3Logger(CustomLogger): - """Test suite for S3 Logger""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.logged_standard_logging_payload = kwargs.get("standard_logging_object") - - -@pytest.mark.asyncio -async def test_dd_llms_obs_redaction(mock_env_vars): - # init DD with turn_off_message_logging=True - litellm._turn_on_debug() - from litellm.types.utils import LiteLLMCommonStrings - - litellm.datadog_llm_observability_params = DatadogLLMObsInitParams( - turn_off_message_logging=True - ) - dd_llms_obs_logger = TestDataDogLLMObsLoggerForRedaction() - test_s3_logger = TestS3Logger() - litellm.callbacks = [dd_llms_obs_logger, test_s3_logger] - - # call litellm - await litellm.acompletion( - model="gpt-4o", - mock_response="Hi there!", - messages=[{"role": "user", "content": "Hello, world!"}], - ) - - # sleep 1 second for logging to complete - await asyncio.sleep(1) - - ################# - # test validation - # 1. both loggers logged a standard_logging_payload - # 2. DD LLM Obs standard_logging_payload has messages and response redacted - # 3. S3 standard_logging_payload does not have messages and response redacted - - assert dd_llms_obs_logger.logged_standard_logging_payload is not None - assert test_s3_logger.logged_standard_logging_payload is not None - - assert ( - dd_llms_obs_logger.logged_standard_logging_payload["messages"][0]["content"] - == "redacted-by-litellm" - ) - assert ( - dd_llms_obs_logger.logged_standard_logging_payload["response"]["choices"][0][ - "message" - ]["content"] - == "redacted-by-litellm" - ) - - assert test_s3_logger.logged_standard_logging_payload["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - assert ( - test_s3_logger.logged_standard_logging_payload["response"]["choices"][0][ - "message" - ]["content"] - == "Hi there!" - ) - - -@pytest.fixture -def mock_env_vars(): - """Mock environment variables for DataDog""" - with patch.dict( - os.environ, {"DD_API_KEY": "test_api_key", "DD_SITE": "us5.datadoghq.com"} - ): - yield - - -@pytest.mark.asyncio -async def test_create_llm_obs_payload(mock_env_vars): - datadog_llm_obs_logger = DataDogLLMObsLogger() - standard_logging_payload = create_standard_logging_payload_with_cache() - payload = datadog_llm_obs_logger.create_llm_obs_payload( - kwargs={ - "model": "gpt-4", - "messages": [{"role": "user", "content": "Hello"}], - "standard_logging_object": standard_logging_payload, - }, - start_time=datetime.now(), - end_time=datetime.now() + timedelta(seconds=1), - ) - - assert payload["name"] == "litellm_llm_call" - assert payload["meta"]["kind"] == "llm" - assert payload["meta"]["input"]["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - assert payload["meta"]["output"]["messages"][0]["content"] == "Hi there!" - assert payload["metrics"]["input_tokens"] == 10 - assert payload["metrics"]["output_tokens"] == 20 - assert payload["metrics"]["total_tokens"] == 30 - - -def create_standard_logging_payload_with_latency_metrics() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object with latency metrics for testing""" - guardrail_info = StandardLoggingGuardrailInformation( - guardrail_name="test_guardrail", - guardrail_status="success", - start_time=1234567890.0, - end_time=1234567890.5, - duration=0.5, # 500ms - guardrail_request={"input": "test input message", "user_id": "test_user"}, - guardrail_response={ - "output": "filtered output", - "flagged": False, - "score": 0.1, - }, - ) - - hidden_params = StandardLoggingHiddenParams( - model_id="model-123", - cache_key="test-cache-key", - api_base="https://api.openai.com", - response_cost="0.05", - litellm_overhead_time_ms=150.0, # 150ms - additional_headers=None, - ) - - return StandardLoggingPayload( - id="test-request-id-latency", - call_type="completion", - response_cost=0.05, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=10, - completion_tokens=20, - startTime=1234567890.0, - endTime=1234567892.0, - completionStartTime=1234567890.8, # 800ms after start - response_time=2.0, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4", model_map_value=None - ), - model="gpt-4", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - error_information=None, - model_parameters={"stream": True}, - hidden_params=hidden_params, - guardrail_information=[guardrail_info], - trace_id="test-trace-id-latency", - custom_llm_provider="openai", - ) - - -def test_latency_metrics_in_metadata(mock_env_vars): - """Test that time to first token, litellm overhead, and guardrail overhead are included in metadata""" - with ( - patch("litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_latency_metrics() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - # Test the metadata generation directly - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - latency_metadata = metadata.get("latency_metrics", {}) - - # Verify time to first token is included (800ms) - assert "time_to_first_token_ms" in latency_metadata - assert ( - abs(latency_metadata["time_to_first_token_ms"] - 800.0) < 0.001 - ) # 0.8 seconds * 1000 with tolerance for floating-point precision - - # Verify litellm overhead is included (150ms) - assert "litellm_overhead_time_ms" in latency_metadata - assert latency_metadata["litellm_overhead_time_ms"] == 150.0 - - # Verify guardrail overhead is included (500ms) - assert "guardrail_overhead_time_ms" in latency_metadata - assert ( - latency_metadata["guardrail_overhead_time_ms"] == 500.0 - ) # 0.5 seconds * 1000 - - # Verify these metrics are also included in the full payload - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - payload_metadata_latency = payload["meta"]["metadata"]["latency_metrics"] - - assert abs(payload_metadata_latency["time_to_first_token_ms"] - 800.0) < 0.001 - assert payload_metadata_latency["litellm_overhead_time_ms"] == 150.0 - assert payload_metadata_latency["guardrail_overhead_time_ms"] == 500.0 - - -def test_latency_metrics_edge_cases(mock_env_vars): - """Test latency metrics with edge cases (missing fields, zero values, etc.)""" - with ( - patch("litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test case 1: No latency metrics present - standard_payload = create_standard_logging_payload_with_cache() - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - - # Should not have latency fields if data is missing/zero - assert "time_to_first_token_ms" not in metadata # Will be 0, so not included - assert ( - "litellm_overhead_time_ms" not in metadata - ) # Not present in hidden_params - assert "guardrail_overhead_time_ms" not in metadata # No guardrail_information - - # Test case 2: Zero time to first token should not be included - standard_payload = create_standard_logging_payload_with_cache() - standard_payload["startTime"] = 1000.0 - standard_payload["completionStartTime"] = 1000.0 # Same time = 0 difference - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - assert "time_to_first_token_ms" not in metadata - - # Test case 3: Missing guardrail duration should not crash - standard_payload = create_standard_logging_payload_with_cache() - standard_payload["guardrail_information"] = [ - StandardLoggingGuardrailInformation( - guardrail_name="test", - guardrail_status="success", - # duration is missing - ) - ] - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - assert "guardrail_overhead_time_ms" not in metadata - - -def test_guardrail_information_in_metadata(mock_env_vars): - """Test that guardrail_information is included in metadata with input/output fields""" - with ( - patch("litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Create a standard payload with guardrail information - standard_payload = create_standard_logging_payload_with_latency_metrics() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - # Create the payload and verify guardrail_information is in metadata - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - metadata = payload["meta"]["metadata"] - - # Verify guardrail_information is present in metadata - assert "guardrail_information" in metadata - assert metadata["guardrail_information"] is not None - - # Verify the guardrail information structure - guardrail_info = metadata["guardrail_information"] - assert guardrail_info[0]["guardrail_name"] == "test_guardrail" - assert guardrail_info[0]["guardrail_status"] == "success" - assert guardrail_info[0]["duration"] == 0.5 - - # Verify input/output fields are present - assert "guardrail_request" in guardrail_info[0] - assert "guardrail_response" in guardrail_info[0] - - # Validate the input/output content - assert guardrail_info[0]["guardrail_request"]["input"] == "test input message" - assert guardrail_info[0]["guardrail_request"]["user_id"] == "test_user" - assert guardrail_info[0]["guardrail_response"]["output"] == "filtered output" - assert guardrail_info[0]["guardrail_response"]["flagged"] is False - assert guardrail_info[0]["guardrail_response"]["score"] == 0.1 - - -def create_standard_logging_payload_with_tool_calls() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object with tool calls for testing""" - return { - "id": "test-request-id-tool-calls", - "trace_id": "test-trace-id-tool-calls", - "call_type": "completion", - "stream": None, - "response_cost": 0.05, - "response_cost_failure_debug_info": None, - "status": "success", - "custom_llm_provider": "openai", - "total_tokens": 50, - "prompt_tokens": 20, - "completion_tokens": 30, - "startTime": 1234567890.0, - "endTime": 1234567891.0, - "completionStartTime": 1234567890.5, - "response_time": 1.0, - "model_map_information": {"model_map_key": "gpt-4", "model_map_value": None}, - "model": "gpt-4", - "model_id": "model-123", - "model_group": "openai-gpt", - "api_base": "https://api.openai.com", - "metadata": { - "user_api_key_hash": "test_hash", - "user_api_key_org_id": None, - "user_api_key_alias": "test_alias", - "user_api_key_team_id": "test_team", - "user_api_key_user_id": "test_user", - "user_api_key_team_alias": "test_team_alias", - "user_api_key_user_email": None, - "user_api_key_end_user_id": None, - "user_api_key_request_route": None, - "spend_logs_metadata": None, - "requester_ip_address": "127.0.0.1", - "requester_metadata": None, - "requester_custom_headers": None, - "prompt_management_metadata": None, - "mcp_tool_call_metadata": None, - "vector_store_request_metadata": None, - "applied_guardrails": None, - "usage_object": None, - "cold_storage_object_key": None, - }, - "cache_hit": False, - "cache_key": None, - "saved_cache_cost": 0.0, - "request_tags": [], - "end_user": None, - "requester_ip_address": "127.0.0.1", - "messages": [ - {"role": "user", "content": "What's the weather?"}, - { - "role": "assistant", - "content": "I'll check the weather for you.", - "tool_calls": [ - { - "id": "call_123", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"location": "NYC"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_123", - "content": '{"temperature": 72, "condition": "sunny"}', - }, - ], - "response": { - "choices": [ - { - "message": { - "role": "assistant", - "content": "It's 72°F and sunny in NYC!", - "tool_calls": [ - { - "id": "call_456", - "type": "function", - "function": { - "name": "format_response", - "arguments": '{"temp": 72, "condition": "sunny"}', - }, - } - ], - } - } - ] - }, - "error_str": None, - "error_information": None, - "model_parameters": {"temperature": 0.7}, - "hidden_params": { - "model_id": "model-123", - "cache_key": None, - "api_base": "https://api.openai.com", - "response_cost": "0.05", - "litellm_overhead_time_ms": None, - "additional_headers": None, - "batch_models": None, - "litellm_model_name": None, - "usage_object": None, - }, - "guardrail_information": None, - "standard_built_in_tools_params": None, - } # type: ignore - - -class TestDataDogLLMObsLoggerToolCalls: - """Simple test suite for DataDog LLM Observability Logger tool call handling""" - - @pytest.fixture - def mock_env_vars(self): - """Mock environment variables for DataDog""" - with patch.dict( - os.environ, {"DD_API_KEY": "test_api_key", "DD_SITE": "us5.datadoghq.com"} - ): - yield - - def test_tool_call_span_kind_mapping(self, mock_env_vars): - """Test that tool call operations are correctly mapped to 'tool' span kind""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test MCP tool call mapping - from litellm.types.utils import CallTypes - - assert ( - logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, "123") - == "tool" - ) - - def test_tool_call_payload_creation(self, mock_env_vars): - """Test that tool call payloads are created correctly""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_tool_calls() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Verify basic payload structure - assert payload.get("name") == "litellm_llm_call" - assert payload.get("status") == "ok" - assert ( - payload.get("meta", {}).get("kind") == "llm" - ) # Regular completion, not tool call - - # Verify metrics - metrics = payload.get("metrics", {}) - assert metrics.get("input_tokens") == 20 - assert metrics.get("output_tokens") == 30 - assert metrics.get("total_tokens") == 50 - - def test_tool_call_messages_preserved(self, mock_env_vars): - """Test that tool call messages are preserved in the payload""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_tool_calls() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Verify input messages include tool calls - meta = payload.get("meta", {}) - input_meta = meta.get("input", {}) - input_messages = input_meta.get("messages", []) - assert len(input_messages) == 3 - - # Check assistant message has tool calls - assistant_msg = input_messages[1] - assert assistant_msg.get("role") == "assistant" - assert "tool_calls" in assistant_msg - tool_calls = assistant_msg.get("tool_calls", []) - assert len(tool_calls) == 1 - tool_call = tool_calls[0] - function_info = tool_call.get("function", {}) - assert function_info.get("name") == "get_weather" - - # Check tool message - tool_msg = input_messages[2] - assert tool_msg.get("role") == "tool" - assert tool_msg.get("tool_call_id") == "call_123" - - def test_tool_call_response_handling(self, mock_env_vars): - """Test that tool calls in response are handled correctly""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_tool_calls() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Verify output messages include tool calls - meta = payload.get("meta", {}) - output_meta = meta.get("output", {}) - output_messages = output_meta.get("messages", []) - assert len(output_messages) == 1 - - output_msg = output_messages[0] - assert output_msg.get("role") == "assistant" - assert "tool_calls" in output_msg - output_tool_calls = output_msg.get("tool_calls", []) - assert len(output_tool_calls) == 1 - output_function_info = output_tool_calls[0].get("function", {}) - assert output_function_info.get("name") == "format_response" - - -def create_standard_logging_payload_with_spend_metrics() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object with spend metrics for testing""" - from datetime import datetime, timezone - - # Create a budget reset time 10 days from now (using "10d" format) - budget_reset_at = datetime.now(timezone.utc) + timedelta(days=10) - - return { - "id": "test-request-id-spend", - "trace_id": "test-trace-id-spend", - "call_type": "completion", - "stream": None, - "response_cost": 0.15, - "response_cost_failure_debug_info": None, - "status": "success", - "custom_llm_provider": "openai", - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - "startTime": 1234567890.0, - "endTime": 1234567891.0, - "completionStartTime": 1234567890.5, - "response_time": 1.0, - "model_map_information": {"model_map_key": "gpt-4", "model_map_value": None}, - "model": "gpt-4", - "model_id": "model-123", - "model_group": "openai-gpt", - "api_base": "https://api.openai.com", - "metadata": { - "user_api_key_hash": "test_hash", - "user_api_key_org_id": None, - "user_api_key_alias": "test_alias", - "user_api_key_team_id": "test_team", - "user_api_key_user_id": "test_user", - "user_api_key_team_alias": "test_team_alias", - "user_api_key_user_email": None, - "user_api_key_end_user_id": None, - "user_api_key_request_route": None, - "user_api_key_spend": 0.67, - "user_api_key_max_budget": 10.0, # $10 max budget - "user_api_key_budget_reset_at": budget_reset_at.isoformat(), # ISO format: 2025-09-26T... - "spend_logs_metadata": None, - "requester_ip_address": "127.0.0.1", - "requester_metadata": None, - "requester_custom_headers": None, - "prompt_management_metadata": None, - "mcp_tool_call_metadata": None, - "vector_store_request_metadata": None, - "applied_guardrails": None, - "usage_object": None, - "cold_storage_object_key": None, - }, - "cache_hit": False, - "cache_key": None, - "saved_cache_cost": 0.0, - "request_tags": [], - "end_user": None, - "requester_ip_address": "127.0.0.1", - "messages": [{"role": "user", "content": "Hello, world!"}], - "response": {"choices": [{"message": {"content": "Hi there!"}}]}, - "error_str": None, - "error_information": None, - "model_parameters": {"stream": False}, - "hidden_params": { - "model_id": "model-123", - "cache_key": None, - "api_base": "https://api.openai.com", - "response_cost": "0.15", - "litellm_overhead_time_ms": None, - "additional_headers": None, - "batch_models": None, - "litellm_model_name": None, - "usage_object": None, - }, - "guardrail_information": None, - "standard_built_in_tools_params": None, - } # type: ignore - - -@pytest.mark.asyncio -async def test_datadog_llm_obs_spend_metrics(mock_env_vars): - """Test that budget metrics are properly extracted and logged""" - datadog_llm_obs_logger = DataDogLLMObsLogger() - - # Create a standard logging payload with spend metrics - payload = create_standard_logging_payload_with_spend_metrics() - - # Show the budget reset time in ISO format - budget_reset_iso = payload["metadata"]["user_api_key_budget_reset_at"] - print(f"Budget reset time (ISO format): {budget_reset_iso}") - from datetime import datetime, timezone - - print(f"Current time: {datetime.now(timezone.utc).isoformat()}") - - # Test the _get_spend_metrics method - spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload) - - # Verify budget metrics are present - assert "user_api_key_max_budget" in spend_metrics - assert spend_metrics["user_api_key_max_budget"] == 10.0 - - assert "user_api_key_budget_reset_at" in spend_metrics - # The budget reset should be a datetime string in ISO format - budget_reset = spend_metrics["user_api_key_budget_reset_at"] - assert isinstance(budget_reset, str) - print(f"Budget reset datetime: {budget_reset}") - # Should be close to 10 days from now - budget_reset_dt = datetime.fromisoformat(budget_reset.replace("Z", "+00:00")) - now = datetime.now(timezone.utc) - time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days - assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days - - print(f"Spend metrics: {spend_metrics}") - - -@pytest.mark.asyncio -async def test_datadog_llm_obs_spend_metrics_no_budget(mock_env_vars): - """Test that spend metrics work when no budget is set""" - datadog_llm_obs_logger = DataDogLLMObsLogger() - - # Create a standard logging payload without budget metadata - payload = create_standard_logging_payload_with_spend_metrics() - - # Remove budget-related metadata to test no-budget scenario - payload["metadata"].pop("user_api_key_max_budget", None) - payload["metadata"].pop("user_api_key_budget_reset_at", None) - - # Test the _get_spend_metrics method - spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload) - - # Verify only response cost is present - assert "response_cost" in spend_metrics - assert spend_metrics["response_cost"] == 0.15 - - # Budget metrics should not be present - assert "user_api_key_max_budget" not in spend_metrics - assert "user_api_key_budget_reset_at" not in spend_metrics - - print(f"Spend metrics (no budget): {spend_metrics}") - - -@pytest.mark.asyncio -async def test_spend_metrics_in_datadog_payload(mock_env_vars): - """Test that spend metrics are correctly included in DataDog LLM Observability payloads""" - from datetime import datetime - - datadog_llm_obs_logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_spend_metrics() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = datadog_llm_obs_logger.create_llm_obs_payload( - kwargs, start_time, end_time - ) - - # Verify basic payload structure - assert payload.get("name") == "litellm_llm_call" - assert payload.get("status") == "ok" - - # Verify spend metrics are included in metadata - meta = payload.get("meta", {}) - assert meta is not None, "Meta section should exist in payload" - - metadata = meta.get("metadata", {}) - assert metadata is not None, "Metadata section should exist in meta" - - spend_metrics = metadata.get("spend_metrics", {}) - assert spend_metrics, "Spend metrics should exist in metadata" - - # Check that all metrics are present - assert "response_cost" in spend_metrics - assert "user_api_key_spend" in spend_metrics - assert "user_api_key_max_budget" in spend_metrics - assert "user_api_key_budget_reset_at" in spend_metrics - - # Verify the values are correct - assert spend_metrics["response_cost"] == 0.15 # response_cost - assert spend_metrics["user_api_key_spend"] == 0.67 # lol - assert spend_metrics["user_api_key_max_budget"] == 10.0 # max budget - - # Verify budget reset is a datetime string in ISO format - budget_reset = spend_metrics["user_api_key_budget_reset_at"] - assert isinstance(budget_reset, str) - print( - f"Budget reset in payload: {budget_reset}" - ) # In StandardLoggingUserAPIKeyMetadata - user_api_key_budget_reset_at: Optional[str] = None - - # In DDLLMObsSpendMetrics - user_api_key_budget_reset_at: str - # Should be close to 10 days from now - from datetime import datetime, timezone - - budget_reset_dt = datetime.fromisoformat(budget_reset.replace("Z", "+00:00")) - now = datetime.now(timezone.utc) - time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days - assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py deleted file mode 100644 index 1043c26c6ec..00000000000 --- a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Tests for the guardrail_translation_mappings registry. - -Validates: -- allm_passthrough_route is registered in the mappings (regression: this was the bug) -""" - -from litellm.llms.pass_through.guardrail_translation import ( - guardrail_translation_mappings, -) -from litellm.llms.pass_through.guardrail_translation.handler import ( - LlmPassthroughRouteHandler, -) -from litellm.types.utils import CallTypes - - -class TestRegistry: - def test_allm_passthrough_route_registered(self): - """Regression: missing this mapping was the root cause of the bug.""" - assert CallTypes.allm_passthrough_route in guardrail_translation_mappings - - def test_allm_passthrough_route_maps_to_llm_passthrough_route_handler(self): - assert ( - guardrail_translation_mappings[CallTypes.allm_passthrough_route] - is LlmPassthroughRouteHandler - ) - - def test_pass_through_still_registered(self): - from litellm.llms.pass_through.guardrail_translation.handler import ( - PassThroughEndpointHandler, - ) - - assert ( - guardrail_translation_mappings[CallTypes.pass_through] - is PassThroughEndpointHandler - ) - diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py deleted file mode 100644 index 020b5de0a2a..00000000000 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py +++ /dev/null @@ -1,296 +0,0 @@ -""" -Test reasoning content preservation in Responses API transformation -""" - -from unittest.mock import AsyncMock - -from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, -) -from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, -) -from litellm.types.utils import ( - Choices, - Delta, - Message, - ModelResponse, - ModelResponseStream, - StreamingChoices, -) - - -class TestReasoningContentStreaming: - """Test reasoning content preservation during streaming""" - - def test_reasoning_content_in_delta(self): - """Test that reasoning content is preserved in streaming deltas""" - # Setup - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="", - role="assistant", - reasoning_content="Let me think about this problem...", - ), - ) - ], - ) - - mock_stream = AsyncMock() - - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=mock_stream, - request_input="Test input", - responses_api_request={}, - ) - - # Execute - transformed_chunk = ( - iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - ) - - # Assert - assert transformed_chunk.delta == "Let me think about this problem..." - assert transformed_chunk.type == "response.reasoning_summary_text.delta" - - def test_mixed_content_and_reasoning(self): - """Test handling of both content and reasoning content""" - # Setup - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="Here is the answer", - role="assistant", - reasoning_content="First, let me analyze...", - ), - ) - ], - ) - - mock_stream = AsyncMock() - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=mock_stream, - request_input="Test input", - responses_api_request={}, - ) - - # Execute - transformed_chunk = ( - iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - ) - - # Assert - assert transformed_chunk.delta == "First, let me analyze..." - assert transformed_chunk.type == "response.reasoning_summary_text.delta" - - def test_no_reasoning_content(self): - """Test handling when no reasoning content is present""" - # Setup - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="Regular content only", - role="assistant", - ), - ) - ], - ) - - mock_stream = AsyncMock() - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=mock_stream, - request_input="Test input", - responses_api_request={}, - ) - - # Execute - transformed_chunk = ( - iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - ) - - # Assert - assert transformed_chunk.delta == "Regular content only" - assert transformed_chunk.type == "response.output_text.delta" - - -class TestReasoningContentFinalResponse: - """Test reasoning content preservation in final response transformation""" - - def test_reasoning_content_in_final_response(self): - """Test that reasoning content is included in final response""" - # Setup - response = ModelResponse( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="Here is my answer", - role="assistant", - reasoning_content="Let me think step by step about this problem...", - ), - ) - ], - ) - - # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, - ) - - # Assert - assert hasattr(responses_api_response, "output") - assert len(responses_api_response.output) > 0 - - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] - assert len(reasoning_items) > 0, "No reasoning item found in output" - - reasoning_item = reasoning_items[0] - assert ( - reasoning_item.content[0].text - == "Let me think step by step about this problem..." - ) - - def test_no_reasoning_content_in_response(self): - """Test handling when no reasoning content in response""" - # Setup - response = ModelResponse( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="Simple answer", - role="assistant", - ), - ) - ], - ) - - # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, - ) - - # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] - assert ( - len(reasoning_items) == 0 - ), "Should have no reasoning items when no reasoning content present" - - def test_multiple_choices_with_reasoning(self): - """Test handling multiple choices, first with reasoning content""" - # Setup - response = ModelResponse( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="First answer", - role="assistant", - reasoning_content="Reasoning for first answer", - ), - ), - Choices( - finish_reason="stop", - index=1, - message=Message( - content="Second answer", - role="assistant", - reasoning_content="Reasoning for second answer", - ), - ), - ], - ) - - # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, - ) - - # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] - assert len(reasoning_items) == 1, "Should have exactly one reasoning item" - assert reasoning_items[0].content[0].text == "Reasoning for first answer" - - -def test_streaming_chunk_id_raw(): - """Test that streaming chunk IDs are raw (not encoded) to match OpenAI format""" - chunk = ModelResponseStream( - id="chunk-123", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content="Hello", role="assistant"), - ) - ], - ) - - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=AsyncMock(), - request_input="Test input", - responses_api_request={}, - custom_llm_provider="openai", - litellm_metadata={"model_info": {"id": "gpt-4"}}, - ) - - result = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - - # Streaming chunk IDs should be raw (like OpenAI's msg_xxx format) - assert result.item_id == "chunk-123" # Should be raw, not encoded - assert not result.item_id.startswith("resp_") # Should NOT have resp_ prefix diff --git a/tests/test_litellm/test_azure_video_router.py b/tests/test_litellm/test_azure_video_router.py deleted file mode 100644 index e7e2e0a01ea..00000000000 --- a/tests/test_litellm/test_azure_video_router.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -Test suite for Azure video router functionality. -Tests that the router method gets called correctly for Azure video generation. -""" - -import pytest -from unittest.mock import Mock, patch, MagicMock -import litellm - - -class TestAzureVideoRouter: - """Test suite for Azure video router functionality""" - - def setup_method(self): - """Setup test fixtures""" - self.model = "azure/sora-2" - self.prompt = "A beautiful sunset over mountains" - self.seconds = "5" - self.size = "1280x720" - - @patch("litellm.videos.main.base_llm_http_handler") - def test_azure_video_generation_router_call_mock(self, mock_handler): - """Test that Azure video generation calls the router method with mock response""" - # Setup mock response - mock_response = { - "id": "video_123", - "model": "sora-2", - "object": "video", - "status": "processing", - "created_at": 1234567890, - "progress": 0, - } - - # Configure the mock handler - mock_handler.video_generation_handler.return_value = mock_response - - # Call the video generation function with mock response - result = litellm.video_generation( - prompt=self.prompt, - model=self.model, - seconds=self.seconds, - size=self.size, - custom_llm_provider="azure", - mock_response=mock_response, - ) - - # Verify the result is a VideoObject with the expected data - assert result.id == mock_response["id"] - assert result.model == mock_response["model"] - assert result.object == mock_response["object"] - assert result.status == mock_response["status"] - assert result.created_at == mock_response["created_at"] - assert result.progress == mock_response["progress"] From bae58eb4e0aa015d5085264e8b0d9a342f163ce5 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Fri, 31 Jul 2026 13:03:13 -0400 Subject: [PATCH 002/245] fix(anthropic): preserve mid-turn system messages Generated with AI Co-Authored-By: Claude Code --- .../chat/guardrail_translation/handler.py | 209 ++++-- .../adapters/transformation.py | 48 +- .../responses_adapters/transformation.py | 49 +- litellm/types/guardrails.py | 47 +- litellm/types/llms/anthropic.py | 12 + .../test_anthropic_guardrail_handler.py | 699 ++++++++++++++++++ ...al_pass_through_adapters_transformation.py | 218 ++++++ .../context_management/test_compact.py | 54 ++ .../test_responses_adapters_transformation.py | 138 ++++ 9 files changed, 1380 insertions(+), 94 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index a549db94224..82707e741c0 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,6 +13,7 @@ Pattern Overview: """ import json +from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger @@ -24,7 +25,6 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTra from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - openai_messages_without_system, openai_messages_without_tool, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -59,14 +59,10 @@ if TYPE_CHECKING: class AnthropicMessagesHandler(BaseTranslation): - """ - Handler for processing Anthropic messages with guardrails. + """Process Anthropic messages with guardrails. - This class provides methods to: - 1. Process input messages (pre-call hook) - 2. Process output responses (post-call hook) - - Methods can be overridden to customize behavior for different message formats. + In-sequence system entries are untrusted client input. This handler scans and preserves + them through guardrail rewrites; downstream provider handling is out of scope. """ def __init__(self): @@ -279,14 +275,26 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply) skip_tool = effective_skip_tool_message_for_guardrail(guardrail_to_apply) - chat_completion_compatible_request = self._translate_to_openai(data) + # 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. + translation_source = { # mutable-ok: API message payload + key: value for key, value in data.items() if key != "system" + } # mutable-ok: API message payload + chat_completion_compatible_request = self._translate_to_openai(translation_source) structured_messages = cast( List[AllMessageValues], chat_completion_compatible_request.get("messages", []), ) - if skip_system: - structured_messages = openai_messages_without_system(structured_messages) + has_midturn_system_message = any( + str(message.get("role") or "").lower() == "system" for message in structured_messages + ) + hoisted_system_message: AllMessageValues | None = None + if not skip_system: + hoisted_system_message = self._hoisted_top_level_system_message(data) + if hoisted_system_message is not None: + structured_messages.insert(0, hoisted_system_message) if skip_tool: structured_messages = openai_messages_without_tool(structured_messages) @@ -346,7 +354,12 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages( + data, + guardrailed_structured_messages, + hoisted_system_message=hoisted_system_message, + preserve_system_messages=has_midturn_system_message, + ) else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( @@ -359,36 +372,120 @@ class AnthropicMessagesHandler(BaseTranslation): return data - @staticmethod - def _write_back_structured_messages(data: dict, structured_messages: list) -> None: - """Convert compressed structured_messages back to Anthropic format and write to data. + def _hoisted_top_level_system_message( + self, data: dict + ) -> AllMessageValues | None: # mutable-ok: API message payload + """Return the system message produced by translating the top-level prompt.""" + system = data.get("system") + if not system: + return None + probe = self._translate_to_openai( + { # mutable-ok: API message payload + "model": data.get("model") or "", + "messages": [], # mutable-ok: API message payload + "system": system, + } + ) + hoisted = probe.get("messages") or [] # mutable-ok: API message payload + return hoisted[0] if hoisted else None - ``anthropic_messages_pt`` merges every run of consecutive user/tool rows - into a single message, so a turn carrying only tool results and the user - turn that follows it come back fused, and the request the model sees no - longer has the boundaries the client sent. Converting a row at a time - would keep them apart but breaks tool pairing: an assistant row whose - tool results sit outside its own call reads as an orphaned tool call, - and under ``modify_params`` the sanitizer answers it with a synthetic - "tool execution skipped" result and drops the real one. Converting each - assistant row together with the tool rows that answer it, and every - other row on its own, satisfies both. - """ + @staticmethod + def _openai_system_message_to_anthropic( + message: dict[str, Any], + ) -> dict[str, Any] | None: # mutable-ok: API message payload + """Convert an OpenAI system message to the client's Anthropic-shaped entry.""" + content = message.get("content") + if isinstance(content, str): + return ( + {"role": "system", "content": content} if content else None # mutable-ok: API message payload + ) # mutable-ok: API message payload + if not isinstance(content, list): + return None + blocks: list[dict[str, Any]] = [] # mutable-ok: API message payload + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text = block.get("text") + if not isinstance(text, str) or not text: + continue + anthropic_block: dict[str, Any] = { # mutable-ok: API message payload + "type": "text", + "text": text, + } # mutable-ok: API message payload + cache_control = block.get("cache_control") + if cache_control: + anthropic_block["cache_control"] = deepcopy(cache_control) + blocks.append(anthropic_block) + return ( + {"role": "system", "content": blocks} if blocks else None # mutable-ok: API message payload + ) # mutable-ok: API message payload + + @staticmethod + def _is_hoisted_top_level_system(message: Any, hoisted_system_message: Any) -> bool: + """Match the hoisted prompt by identity, or by value after serialization.""" + if hoisted_system_message is None: + return False + if message is hoisted_system_message: + return True + return ( + isinstance(message, dict) and isinstance(hoisted_system_message, dict) and message == hoisted_system_message + ) + + @staticmethod + def _write_back_structured_messages( + data: dict, # mutable-ok: API message payload + structured_messages: list, # mutable-ok: API message payload + hoisted_system_message: Any = None, + preserve_system_messages: bool = False, + ) -> None: + """Write a guardrail's structured-message rewrite back without losing corrections.""" from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, group_tool_exchanges, ) + def _is_system(message: Any) -> bool: + return isinstance(message, dict) and str(message.get("role") or "").lower() == "system" + model = str(data.get("model") or "") - non_system = [m for m in structured_messages if m.get("role") != "system"] - groups = tuple([non_system[index] for index in group] for group in group_tool_exchanges(non_system)) or ( - non_system, - ) - converted = [ - message - for group in groups - for message in anthropic_messages_pt(messages=group, model=model, llm_provider="anthropic") - ] + converted: list = [] # mutable-ok: API message payload + + def _convert_run(run: list) -> None: # mutable-ok: API message payload + for group in group_tool_exchanges(run): + converted.extend( + anthropic_messages_pt( + messages=[ # mutable-ok: API message payload + run[index] for index in group + ], # mutable-ok: API message payload + model=model, + llm_provider="anthropic", + ) + ) + + run: list = [] # mutable-ok: API message payload + hoisted_dropped = False + for message in structured_messages: + if not _is_system(message): + run.append(message) + continue + _convert_run(run) + run = [] # mutable-ok: API message payload + if not hoisted_dropped and AnthropicMessagesHandler._is_hoisted_top_level_system( + message, hoisted_system_message + ): + hoisted_dropped = True + continue + if preserve_system_messages: + anthropic_system = AnthropicMessagesHandler._openai_system_message_to_anthropic(message) + if anthropic_system is not None: + converted.append(anthropic_system) + _convert_run(run) + if not any(not _is_system(message) for message in converted): + converted.extend( + anthropic_messages_pt( + messages=[], model=model, llm_provider="anthropic" + ) # mutable-ok: API message payload + ) # mutable-ok: API message payload for msg in converted: content = msg.get("content") if isinstance(content, list): @@ -397,6 +494,29 @@ class AnthropicMessagesHandler(BaseTranslation): block.pop("cache_control", None) data["messages"] = converted + @staticmethod + def _extract_midturn_system_text( + message: dict[str, Any], # mutable-ok: API message payload + msg_idx: int, + texts_to_check: list[str], # mutable-ok: API message payload + task_mappings: list[tuple[int, int | None]], # mutable-ok: API message payload + ) -> None: + content = message.get("content") + if isinstance(content, str): + if content: + texts_to_check.append(content) + task_mappings.append((msg_idx, None)) + return + if not isinstance(content, list): + return + for content_idx, content_item in enumerate(content): + if not isinstance(content_item, dict) or content_item.get("type") != "text": + continue + text_str = content_item.get("text") + if isinstance(text_str, str) and text_str: + texts_to_check.append(text_str) + task_mappings.append((msg_idx, content_idx)) + def extract_request_tool_names(self, data: dict) -> List[str]: """Extract tool names from Anthropic messages request (tools[].name).""" names: List[str] = [] @@ -415,15 +535,18 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system_message: bool = False, skip_tool_message: bool = False, ) -> None: - """ - Extract text content and images from a message. - - Override this method to customize text/image extraction logic. - """ - role = str(message.get("role") or "").lower() - if skip_system_message and role == "system": + """Extract text content and images from a message.""" + role = str(message.get("role") or "") + if role == "system": + # Match the adapter's filtering so positional guardrail write-back stays aligned. + self._extract_midturn_system_text( + message=message, + msg_idx=msg_idx, + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) return - if skip_tool_message and role == "tool": + if skip_tool_message and role.lower() == "tool": return content = message.get("content", None) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 86c9c1db481..707d53e9006 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -85,12 +85,12 @@ from litellm.llms.anthropic.experimental_pass_through.context_management import ) from litellm.types.llms.anthropic import ( ANTHROPIC_HOSTED_TOOLS, + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, AnthropicFinishReason, AnthropicMessagesRequest, + AnthropicMessagesSystemMessageParam, AnthropicMessagesToolChoice, - AnthropicMessagesUserMessageParam, AnthropicResponseContentBlockRedactedThinking, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, @@ -354,12 +354,7 @@ class LiteLLMAnthropicMessagesAdapter: def translate_anthropic_messages_to_openai( self, - messages: List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + messages: List[AllAnthropicPassThroughMessageValues], # mutable-ok: API message payload model: Optional[str] = None, ) -> List: new_messages: List[AllMessageValues] = [] @@ -367,6 +362,11 @@ class LiteLLMAnthropicMessagesAdapter: user_message: Optional[ChatCompletionUserMessage] = None tool_message_list: List[ChatCompletionToolMessage] = [] new_user_content_list: List[Union[ChatCompletionTextObject, ChatCompletionImageObject]] = [] + if m["role"] == "system": + system_message = self._translate_midturn_system_message_to_openai(m, model) + if system_message is not None: + new_messages.append(system_message) + continue ## USER MESSAGE ## if m["role"] == "user": ## translate user message @@ -867,6 +867,29 @@ class LiteLLMAnthropicMessagesAdapter: for def_schema in schema[key].values(): LiteLLMAnthropicMessagesAdapter._add_additional_properties_false(def_schema) + def _translate_midturn_system_message_to_openai( + self, + message: AnthropicMessagesSystemMessageParam, + model: str | None, + ) -> ChatCompletionSystemMessage | None: + """Translate an in-sequence system entry without changing its role or position.""" + content = message.get("content") + if isinstance(content, str): + return ChatCompletionSystemMessage(role="system", content=content) if content else None + if not isinstance(content, list): + return None + text_parts: list[ChatCompletionTextObject] = [] # mutable-ok: API message payload + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text = block.get("text") + if not text: + continue + text_obj = ChatCompletionTextObject(type="text", text=text) + self._add_cache_control_if_applicable(block, text_obj, model) + text_parts.append(text_obj) + return ChatCompletionSystemMessage(role="system", content=text_parts) if text_parts else None + def _add_system_message_to_messages( self, new_messages: List[AllMessageValues], @@ -1068,13 +1091,8 @@ class LiteLLMAnthropicMessagesAdapter: tool_name_mapping: Dict[str, str] = {} ## CONVERT ANTHROPIC MESSAGES TO OPENAI - messages_list: List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]] = cast( - List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + messages_list = cast( + List[AllAnthropicPassThroughMessageValues], anthropic_message_request["messages"], ) new_messages = self.translate_anthropic_messages_to_openai( diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 172e54de98e..cbe36100eb0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -6,6 +6,7 @@ path used for OpenAI and Azure models. """ import json +from collections.abc import Iterable from typing import Any, Dict, List, Optional, Union, cast from litellm.litellm_core_utils.reasoning_effort_utils import ( @@ -15,15 +16,15 @@ from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, ) from litellm.types.llms.anthropic import ( + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, AnthropicFinishReason, AnthropicMessagesRequest, AnthropicMessagesToolChoice, - AnthropicMessagesUserMessageParam, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockToolUse, + AnthropicSystemMessageContent, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -54,19 +55,32 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return source.get("url") return None + @staticmethod + def _translate_midturn_system_content_to_responses( + content: Union[str, Iterable[AnthropicSystemMessageContent]], + ) -> list[dict[str, str]]: # mutable-ok: API message payload + """Convert in-sequence system content to Responses input-text parts.""" + if isinstance(content, str): + return ( + [{"type": "input_text", "text": content}] if content else [] # mutable-ok: API message payload + ) # mutable-ok: API message payload + if not isinstance(content, list): + return [] # mutable-ok: API message payload + return [ # mutable-ok: API message payload + {"type": "input_text", "text": text} # mutable-ok: API message payload + for block in content + if isinstance(block, dict) and block.get("type") == "text" and (text := block.get("text")) + ] + def translate_messages_to_responses_input( self, - messages: List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + messages: List[AllAnthropicPassThroughMessageValues], # mutable-ok: API message payload ) -> List[Dict[str, Any]]: """ Convert Anthropic messages list to Responses API `input` items. Mapping: + system text -> message(role=system, input_text) user text -> message(role=user, input_text) user image -> message(role=user, input_image) user tool_result -> function_call_output @@ -76,6 +90,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter: input_items: List[Dict[str, Any]] = [] for m in messages: + if m["role"] == "system": + system_parts = self._translate_midturn_system_content_to_responses(m.get("content")) + if system_parts: + input_items.append( + { # mutable-ok: API message payload + "type": "message", + "role": "system", + "content": system_parts, + } + ) + continue + role = m["role"] content = m.get("content") @@ -287,12 +313,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: """ model: str = anthropic_request["model"] messages_list = cast( - List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + List[AllAnthropicPassThroughMessageValues], anthropic_request["messages"], ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index af419d8cb6f..2e0da24ccda 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -11,12 +11,24 @@ from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import ( BlockCodeExecutionGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( + CiscoAIDefenseGuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import ( EnkryptAIGuardrailConfigs, ) from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( GraySwanGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HeadroomGuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( + HiddenlayerGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) @@ -29,38 +41,26 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import ( from litellm.types.proxy.guardrails.guardrail_hooks.promptguard import ( PromptGuardConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( - XecGuardConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, ) from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( - HiddenlayerGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( - QostodianNexusConfigModel, -) from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import ( RepelloAIGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( - VigilGuardGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( - CiscoAIDefenseGuardrailConfigModel, -) from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( SingulrGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( - HeadroomGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( - CompresrGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( + VigilGuardGuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( + XecGuardConfigModel, ) """ @@ -743,7 +743,10 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up "When True, unified guardrails skip system-role messages when building " "evaluation inputs (texts and structured_messages). When False, system " "messages are included even if litellm_settings sets a global skip. When " - "None, use the global litellm.skip_system_message_in_guardrail setting." + "None, use the global litellm.skip_system_message_in_guardrail setting. " + "For Anthropic /v1/messages, the flag applies only to the trusted top-level " + "system prompt. In-sequence system entries are untrusted client input and remain " + "in texts and structured_messages." ), ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index c24d072217a..29faf500b3d 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -365,8 +365,20 @@ class AnthropicSystemMessageContent(TypedDict, total=False): cache_control: Optional[Union[dict, ChatCompletionCachedContent]] +class AnthropicMessagesSystemMessageParam(TypedDict, total=False): + role: Required[Literal["system"]] + content: Required[Union[str, Iterable[AnthropicSystemMessageContent]]] + + AllAnthropicMessageValues = Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam] +# System is not a native Anthropic message role; only pass-through adapters use this union. +AllAnthropicPassThroughMessageValues = Union[ + AnthropicMessagesUserMessageParam, + AnthopicMessagesAssistantMessageParam, + AnthropicMessagesSystemMessageParam, +] + class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): max_tokens: Optional[int] 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 48acdd348e9..7757b0fa5a4 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 @@ -75,6 +75,51 @@ class MockRecordingGuardrail(CustomGuardrail): return inputs +class MockMaskingGuardrail(CustomGuardrail): + """Capture request inputs and mask one known prohibited value.""" + + def __init__(self, skip_system_message_in_guardrail: Optional[bool] = True): + super().__init__(guardrail_name="masking-test") + self.skip_system_message_in_guardrail = skip_system_message_in_guardrail + self.inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.inputs = inputs.copy() + masked_inputs = inputs.copy() + masked_inputs["texts"] = [ + "[MASKED]" if text == "prohibited correction" else text for text in inputs.get("texts", []) + ] + return masked_inputs + + +class MockCompactingGuardrail(CustomGuardrail): + """Stand in for a compaction guardrail that rewrites `structured_messages` wholesale.""" + + def __init__(self, replacement_messages: list): + super().__init__(guardrail_name="compacting-test") + self.replacement_messages = replacement_messages + self.inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.inputs = inputs.copy() + rewritten = inputs.copy() + # A new list object -- this is what signals a rewrite to the handler. + rewritten["structured_messages"] = list(self.replacement_messages) + return rewritten + + class TestAnthropicMessagesHandlerStreamingRequestData: """Post-call guardrails on streaming /v1/messages receive the response and identity metadata""" @@ -210,6 +255,660 @@ class TestAnthropicMessagesHandlerInputProcessing: assert data.get("litellm_metadata", {}).get("guardrails") assert guardrail.dynamic_params == {"policy_id": "policy-123"} + @pytest.mark.asyncio + async def test_midturn_system_correction_is_guardrailed_when_top_level_system_is_skipped( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + { + "role": "system", + "content": [ + {"type": "unsupported", "text": "discarded text"}, + {"type": "text", "text": "prohibited correction"}, + ], + }, + ], + } + + 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 "trusted top-level system prompt" not in guardrail.inputs["texts"] + assert data["messages"][1]["content"][0]["text"] == "discarded text" + assert data["messages"][1]["content"][1]["text"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_string_midturn_system_correction_is_guardrailed(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "system", "content": "prohibited correction"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert guardrail.inputs["texts"] == ["prohibited correction"] + assert data["messages"][0]["content"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_unsupported_midturn_system_content_is_not_guardrailed(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + { + "role": "system", + "content": [{"type": "image", "source": {"type": "url"}}], + } + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is None + + @pytest.mark.asyncio + async def test_skip_system_message_excludes_only_hoisted_top_level_system(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "system", "content": "prohibited correction"}, + {"role": "user", "content": "continue"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + structured = guardrail.inputs["structured_messages"] + assert [m["role"] for m in structured] == ["user", "system", "user"] + assert structured[1]["content"] == "prohibited correction" + + @pytest.mark.asyncio + async def test_default_skip_false_scans_midturn_system_and_hoists_top_level_system( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail(skip_system_message_in_guardrail=None) + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "system", "content": "prohibited correction"}, + ], + } + + 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"] + 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["messages"][1]["content"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included( + self, + ): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail(skip_system_message_in_guardrail=None) + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "system", "content": "prohibited correction"}, + {"role": "user", "content": "latest question"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + 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 + 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 + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None]) + async def test_midturn_system_text_extraction_matches_translation_in_both_skip_modes( + self, + skip_system_message_in_guardrail: Optional[bool], + ): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail(skip_system_message_in_guardrail=skip_system_message_in_guardrail) + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "safe text"}, + { + "role": "system", + "content": [ + {"type": "text", "text": ""}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + {"type": "text", "text": "prohibited correction"}, + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + assert texts == ["safe text", "prohibited correction"] + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + assert data["messages"][1]["content"][2]["text"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_bedrock_masking_slice_stays_aligned_with_midturn_system(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + { + "role": "system", + "content": [ + {"type": "text", "text": "prohibited correction"}, + {"type": "text", "text": "second correction"}, + ], + }, + {"role": "user", "content": "latest question"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + total = sum(bedrock._count_message_texts(m) for m in structured) + assert total == len(texts) + + latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") + assert latest_user_index == 2 + scanned_slice = bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, + ) + assert scanned_slice == (3, 1) + + merged = bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=scanned_slice, + scanned_role_subset=True, + ) + assert merged == [ + "safe text", + "prohibited correction", + "second correction", + "{MASKED}", + ] + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_midturn_system_messages(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + { + "role": "system", + "content": [{"type": "text", "text": "use the corrected result"}], + }, + {"role": "user", "content": "continue"}, + ] + ) + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "continue"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "system", "user"] + assert data["messages"][1]["content"] == [{"type": "text", "text": "use the corrected result"}] + assert data["messages"][0]["content"] == [{"type": "text", "text": "compacted history"}] + assert data["messages"][2]["content"] == [{"type": "text", "text": "continue"}] + assert data["system"] == "trusted top-level system prompt" + + @pytest.mark.asyncio + async def test_compaction_rewrite_does_not_duplicate_hoisted_top_level_system(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "trusted top-level system prompt"}, + {"role": "user", "content": "compacted history"}, + {"role": "system", "content": "use the corrected result"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "system"] + assert data["messages"][1]["content"] == "use the corrected result" + assert data["system"] == "trusted top-level system prompt" + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_leading_midturn_system_when_system_is_skipped( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "compacted history"}, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "original history"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "use the corrected result" + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_leading_correction_when_top_level_system_hoists_nothing( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "compacted history"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}], + "messages": [ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "original history"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "use the corrected result" + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_leading_correction_when_hoisted_prompt_is_dropped( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "CLIENT CORRECTION"}, + {"role": "user", "content": "compacted history"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "TRUSTED", + "messages": [ + {"role": "system", "content": "CLIENT CORRECTION"}, + {"role": "user", "content": "original history"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert guardrail.inputs["structured_messages"][0] == { + "role": "system", + "content": "TRUSTED", + } + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "CLIENT CORRECTION" + assert data["system"] == "TRUSTED" + + @pytest.mark.asyncio + async def test_compaction_rewrite_drops_hoisted_prompt_matched_by_content_copy(self): + import json + + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + json.loads(json.dumps({"role": "system", "content": "TRUSTED"})), + {"role": "user", "content": "compacted history"}, + {"role": "system", "content": "CLIENT CORRECTION"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "TRUSTED", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "CLIENT CORRECTION"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "system"] + assert data["messages"][1]["content"] == "CLIENT CORRECTION" + assert data["system"] == "TRUSTED" + + @pytest.mark.asyncio + async def test_compaction_rewrite_preserves_cache_control_on_system_blocks(self): + """ + `cache_control` on an in-sequence system text block survives the write-back, and is + copied rather than aliased into the guardrail's own returned list. + """ + handler = AnthropicMessagesHandler() + source_cache_control = {"type": "ephemeral"} + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + { + "role": "system", + "content": [ + { + "type": "text", + "text": "use the corrected result", + "cache_control": source_cache_control, + } + ], + }, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"] == [ + { + "type": "text", + "text": "use the corrected result", + "cache_control": {"type": "ephemeral"}, + } + ] + assert data["messages"][1]["content"][0]["cache_control"] is not source_cache_control + + @pytest.mark.asyncio + async def test_compaction_rewrite_rstrips_trailing_assistant_in_each_run(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + {"role": "assistant", "content": "earlier "}, + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "continue"}, + {"role": "assistant", "content": "prefill "}, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == [ + "user", + "assistant", + "system", + "user", + "assistant", + ] + assert data["messages"][1]["content"] == [{"type": "text", "text": "earlier"}] + assert data["messages"][-1]["content"] == [{"type": "text", "text": "prefill"}] + + @pytest.mark.asyncio + async def test_compaction_rewrite_drops_text_free_system_message(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + {"role": "system", "content": [{"type": "text", "text": ""}]}, + {"role": "system", "content": ""}, + {"role": "user", "content": "continue"}, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "continue"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "user"] + assert data["messages"][0]["content"] == [{"type": "text", "text": "compacted history"}] + assert data["messages"][1]["content"] == [{"type": "text", "text": "continue"}] + + @pytest.mark.asyncio + async def test_noncanonical_system_role_casing_is_still_scanned(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "System", "content": "prohibited correction"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert "prohibited correction" in guardrail.inputs["texts"] + assert data["messages"][1]["content"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_tool_result_turns_have_a_preexisting_alignment_gap(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + tool_loop = [ + {"role": "user", "content": "call the tool"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu_1", "name": "get", "input": {"a": 1}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [{"type": "text", "text": "tool output"}], + } + ], + }, + ] + + async def _slice_for(messages: list): + guardrail = MockMaskingGuardrail() + data = {"model": "claude-3-5-sonnet-20241022", "messages": messages} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + target_index = bedrock._find_latest_message_index(structured, target_role="user") + return ( + sum(bedrock._count_message_texts(m) for m in structured) - len(texts), + bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=target_index, + texts=texts, + ), + ) + + with_system = await _slice_for( + tool_loop + + [ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "latest question"}, + ] + ) + without_system = await _slice_for(tool_loop + [{"role": "user", "content": "latest question"}]) + + assert with_system == without_system == (1, None) + + @pytest.mark.asyncio + async def test_compaction_rewrite_to_only_system_messages_is_rejected(self): + import litellm + + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[{"role": "system", "content": "use the corrected result"}] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + with patch.object(litellm, "modify_params", False): + with pytest.raises(litellm.BadRequestError, match="at least one non-system message"): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + @pytest.mark.asyncio + async def test_compaction_rewrite_to_only_system_messages_repaired_with_modify_params( + self, + ): + import litellm + + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[{"role": "system", "content": "use the corrected result"}] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + with patch.object(litellm, "modify_params", True): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "use the corrected result" + + @pytest.mark.asyncio + async def test_compaction_rewrite_without_system_messages_is_unchanged(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail(replacement_messages=[{"role": "user", "content": "compacted history"}]) + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + {"role": "user", "content": "c"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"] == [{"role": "user", "content": [{"type": "text", "text": "compacted history"}]}] + @pytest.mark.asyncio async def test_process_output_streaming_response_empty_choices(self): """Test that streaming response with empty choices doesn't raise IndexError diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index c0c6e315b5b..b72620f9918 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -413,6 +413,224 @@ def test_translate_anthropic_messages_to_openai_tool_message_placement(): ), "Tool message should be placed before user message" +@pytest.mark.parametrize( + ("system_content", "expected_content"), + [ + ("Use the corrected result.", "Use the corrected result."), + ( + [{"type": "text", "text": "Use the corrected result."}], + [{"type": "text", "text": "Use the corrected result."}], + ), + ( + [ + { + "type": "image", + "source": {"type": "url", "url": "https://example.com/a.png"}, + }, + {"type": "text", "text": "Use the corrected result."}, + ], + [{"type": "text", "text": "Use the corrected result."}], + ), + ( + [ + {"type": "text", "text": "First correction."}, + {"type": "text", "text": "Second correction."}, + ], + [ + {"type": "text", "text": "First correction."}, + {"type": "text", "text": "Second correction."}, + ], + ), + ], +) +def test_translate_anthropic_messages_to_openai_preserves_midturn_system_correction( + system_content: object, + expected_content: object, +): + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01234", + "name": "get_weather", + "input": {"location": "Boston"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01234", + "content": "Rainy, 55°F", + } + ], + }, + {"role": "system", "content": system_content}, + {"role": "user", "content": "Continue."}, + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="claude-3-5-sonnet-20240620", + ) + + assert result == [ + { + "role": "assistant", + "content": None, + "thinking_blocks": None, + "tool_calls": [ + { + "id": "toolu_01234", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "toolu_01234", + "content": "Rainy, 55°F", + }, + {"role": "system", "content": expected_content}, + {"role": "user", "content": "Continue."}, + ] + + +def test_translate_anthropic_messages_to_openai_preserves_midturn_system_cache_control(): + """ + `cache_control` on an in-sequence system text block survives, matching how the + hoisted top-level `system` prompt and user text blocks are already handled. + """ + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Use the corrected result.", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="claude-3-5-sonnet-20240620", + ) + + assert result == [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Use the corrected result.", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + +def test_translate_anthropic_messages_to_openai_drops_midturn_system_cache_control_for_non_claude(): + """ + `cache_control` goes through the same `_add_cache_control_if_applicable` gate as the + hoisted top-level prompt and user text blocks, so a non-Claude *requested model name* + does not get it. That gate is a best-effort check of the requested name before routing + (behind the proxy it is often a public alias), not a guarantee about the backend that + ultimately serves the request. + """ + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Use the corrected result.", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="gpt-4o", + ) + + assert result == [ + { + "role": "system", + "content": [{"type": "text", "text": "Use the corrected result."}], + } + ] + + +@pytest.mark.parametrize( + "system_content", + [ + "", + [{"type": "text", "text": ""}], + [ + { + "type": "image", + "source": {"type": "url", "url": "https://example.com/a.png"}, + } + ], + None, + ], +) +def test_translate_anthropic_messages_to_openai_drops_empty_midturn_system( + system_content: object, +): + messages = [{"role": "system", "content": system_content}] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="claude-3-5-sonnet-20240620", + ) + + assert result == [] + + +def test_translate_anthropic_to_openai_orders_top_level_and_midturn_system(): + """ + Request level: the trusted top-level prompt is hoisted to index 0 exactly once and the + in-sequence correction keeps its own position and `role: "system"` -- no duplication of + either, and no reordering of the surrounding turns. + """ + openai_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request={ + "model": "claude-3-5-sonnet-20240620", + "max_tokens": 100, + "system": "Trusted top-level prompt.", + "messages": [ + {"role": "user", "content": "First question."}, + {"role": "assistant", "content": "First answer."}, + {"role": "system", "content": "Use the corrected result."}, + {"role": "user", "content": "Continue."}, + ], + } + ) + + assert openai_request["messages"] == [ + {"role": "system", "content": "Trusted top-level prompt."}, + {"role": "user", "content": "First question."}, + {"role": "assistant", "content": "First answer.", "thinking_blocks": None}, + {"role": "system", "content": "Use the corrected result."}, + {"role": "user", "content": "Continue."}, + ] + + def test_translate_openai_content_to_anthropic_empty_function_arguments(): """Test that empty function arguments are handled safely and don't cause JSON parsing errors.""" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 9c8df1c79f9..6cc1d9e5add 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -2475,3 +2475,57 @@ def test_endpoint_runs_failure_hook_on_500_context_management_error(): body = response.json() assert body["type"] == "error" failure_hook.assert_awaited_once() + + +def test_count_effective_tokens_counts_midturn_system_correction(): + from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import ( + _count_effective_tokens, + ) + + base: List[Dict[str, Any]] = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + correction = { + "role": "system", + "content": [{"type": "text", "text": "use the corrected result " * 20}], + } + + without_correction = _count_effective_tokens( + model=MODEL, effective_messages=base, compaction_block=None, tools=None + ) + with_correction = _count_effective_tokens( + model=MODEL, + effective_messages=base + [correction], + compaction_block=None, + tools=None, + ) + + assert with_correction > without_correction + + +def test_build_summary_messages_keeps_midturn_system_correction_in_place(): + from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import ( + _build_summary_messages, + ) + + summary_messages = _build_summary_messages( + effective_messages=[ + {"role": "user", "content": "original question"}, + {"role": "system", "content": "use the corrected result"}, + {"role": "assistant", "content": "acknowledged"}, + ], + prompt="summarize the conversation", + system="caller system prompt", + ) + + assert [m["role"] for m in summary_messages] == [ + "system", + "user", + "system", + "assistant", + "user", + ] + assert summary_messages[0]["content"] == "caller system prompt" + assert summary_messages[2]["content"] == "use the corrected result" + assert summary_messages[-1]["content"] == "summarize the conversation" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 606ff39b35e..8963012ecd5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -9,6 +9,8 @@ import sys from typing import Any, Dict, List from unittest.mock import MagicMock +import pytest + sys.path.insert(0, os.path.abspath("../../../../../../..")) from litellm.constants import ( @@ -221,6 +223,106 @@ class TestTranslateMessagesToResponsesInput: {"type": "input_text", "text": "Second part."}, ] + @pytest.mark.parametrize( + "system_content", + [ + "Use the corrected result.", + [{"type": "text", "text": "Use the corrected result."}], + [ + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + {"type": "text", "text": "Use the corrected result."}, + ], + ], + ) + def test_midturn_system_correction_stays_system_in_sequence(self, system_content: object): + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01234", + "name": "get_weather", + "input": {"location": "Boston"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01234", + "content": "Rainy, 55°F", + } + ], + }, + {"role": "system", "content": system_content}, + {"role": "user", "content": "Continue."}, + ] + + result = _translate_messages(messages) + + assert result == [ + { + "type": "function_call", + "call_id": "toolu_01234", + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + { + "type": "function_call_output", + "call_id": "toolu_01234", + "output": "Rainy, 55°F", + }, + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": "Use the corrected result."}], + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + }, + ] + + def test_midturn_system_correction_keeps_multiple_text_blocks(self): + messages = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "First correction."}, + {"type": "text", "text": "Second correction."}, + ], + } + ] + + assert _translate_messages(messages) == [ + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "First correction."}, + {"type": "input_text", "text": "Second correction."}, + ], + } + ] + + @pytest.mark.parametrize( + "system_content", + [ + "", + [{"type": "text", "text": ""}], + [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}], + None, + ], + ) + def test_empty_or_unsupported_midturn_system_correction_is_dropped(self, system_content: object): + messages = [{"role": "system", "content": system_content}] + + assert _translate_messages(messages) == [] + def test_user_base64_image(self): """User message with base64 image source becomes input_image with data URL.""" messages = [ @@ -722,6 +824,42 @@ class TestTranslateRequestBroaderCoverage: kwargs = _ADAPTER.translate_request(req) assert kwargs["instructions"] == "You are a helpful assistant." + def test_top_level_system_and_midturn_correction_are_not_duplicated(self): + """ + Request level: the trusted top-level prompt goes to `instructions` only, and the + in-sequence correction stays a `role: "system"` input item in its original position. + Neither appears twice, and the surrounding turns keep their order. + """ + req = _make_request( + system="Trusted top-level prompt.", + messages=[ + {"role": "user", "content": "First question."}, + {"role": "system", "content": "Use the corrected result."}, + {"role": "user", "content": "Continue."}, + ], + ) + + kwargs = _ADAPTER.translate_request(req) + + assert kwargs["instructions"] == "Trusted top-level prompt." + assert kwargs["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "First question."}], + }, + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": "Use the corrected result."}], + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + }, + ] + def test_system_list_of_text_blocks_joined(self): req = _make_request( system=[ From 81a80b8c632a9b5d942193bb71eda601630e9906 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:02:27 +0000 Subject: [PATCH 003/245] fix(anthropic): preserve speed=fast in usage for /v1/messages and pass-through Fast mode is priced with a provider-specific multiplier applied off usage.speed, but only chat completions kept that field. The Messages route rebuilt usage with empty optional params, stream reassembly dropped speed and inference_geo, and the pass-through handler never read speed off the request body, so fast-mode spend was logged at the standard rate. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- .../streaming_chunk_builder_utils.py | 16 +++ .../anthropic_passthrough_logging_handler.py | 24 +++- .../streaming_chunk_builder_utils.py | 2 + .../test_litellm_logging.py | 35 ++++++ .../test_streaming_chunk_builder_utils.py | 39 ++++++ ...t_anthropic_passthrough_logging_handler.py | 113 ++++++++++++++++++ 7 files changed, 228 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c2dc7189934..15330a2a910 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3278,7 +3278,7 @@ class Logging(LiteLLMLoggingBaseClass): model=self.model, messages=[], logging_obj=self, - optional_params={}, + optional_params=self.optional_params or {}, api_key="", request_data={}, encoding=litellm.encoding, diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d52d9849310..5ceb64547a0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -569,6 +569,8 @@ class ChunkProcessor: # lost and 1h cache writes get billed at the 5m rate. cache_creation_token_details: Optional[CacheCreationTokenDetails] = None cost: Optional[float] = None + inference_geo: Optional[str] = None + speed: Optional[str] = None for chunk in chunks: usage_chunk = self._extract_usage_chunk(chunk) @@ -627,6 +629,13 @@ class ChunkProcessor: if usage_chunk_dict["cost"] is not None: cost = usage_chunk_dict["cost"] + chunk_inference_geo = getattr(usage_chunk, "inference_geo", None) + if isinstance(chunk_inference_geo, str): + inference_geo = chunk_inference_geo + chunk_speed = getattr(usage_chunk, "speed", None) + if isinstance(chunk_speed, str): + speed = chunk_speed + prompt_tokens_details = self._attach_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details ) @@ -647,6 +656,8 @@ class ChunkProcessor: completion_tokens_details=completion_tokens_details, prompt_tokens_details=prompt_tokens_details, cost=cost, + inference_geo=inference_geo, + speed=speed, ) @staticmethod @@ -806,6 +817,11 @@ class ChunkProcessor: if cost is not None: setattr(returned_usage, "cost", cost) + if calculated_usage_per_chunk["inference_geo"] is not None: + setattr(returned_usage, "inference_geo", calculated_usage_per_chunk["inference_geo"]) + if calculated_usage_per_chunk["speed"] is not None: + setattr(returned_usage, "speed", calculated_usage_per_chunk["speed"]) + # Return a new usage object with the new values returned_usage = Usage(**returned_usage.model_dump()) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 50e90699194..112780581de 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -80,7 +80,9 @@ class AnthropicPassthroughLoggingHandler: model=model, messages=[], logging_obj=logging_obj, - optional_params={}, + optional_params=AnthropicPassthroughLoggingHandler._cost_relevant_request_params( + request_body or kwargs.get("request_body") + ), api_key="", request_data={}, encoding=litellm.encoding, @@ -102,6 +104,15 @@ class AnthropicPassthroughLoggingHandler: "kwargs": kwargs, } + @staticmethod + def _cost_relevant_request_params(request_body: Optional[dict]) -> dict: + """ + Request params that change how the response is priced, and so must reach the + usage-building paths. Anthropic's ``speed=fast`` multiplies non-cache token cost. + """ + speed = (request_body or {}).get("speed") + return {"speed": speed} if isinstance(speed, str) else {} + @staticmethod def _get_user_from_metadata( passthrough_logging_payload: PassthroughStandardLoggingPayload, @@ -315,6 +326,7 @@ class AnthropicPassthroughLoggingHandler: - Logs in litellm callbacks """ + speed = AnthropicPassthroughLoggingHandler._cost_relevant_request_params(request_body).get("speed") model = request_body.get("model", "") # Check if it's available in the logging object if ( @@ -334,6 +346,7 @@ class AnthropicPassthroughLoggingHandler: all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + speed=speed, ) except Exception as e: # stream_chunk_builder re-raises assembly failures (as litellm.APIError) @@ -355,6 +368,7 @@ class AnthropicPassthroughLoggingHandler: complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( all_chunks=all_chunks, model=model, + speed=speed, ) except Exception as e: verbose_proxy_logger.warning( @@ -419,6 +433,7 @@ class AnthropicPassthroughLoggingHandler: all_chunks: Sequence[Union[str, bytes]], litellm_logging_obj: LiteLLMLoggingObj, model: str, + speed: Optional[str] = None, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: """ Builds complete response from raw Anthropic chunks. @@ -443,11 +458,13 @@ class AnthropicPassthroughLoggingHandler: all_chunks=collapsed, litellm_logging_obj=litellm_logging_obj, model=model, + speed=speed, ) return AnthropicPassthroughLoggingHandler._build_complete_streaming_response_legacy( all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + speed=speed, ) # Anthropic SSE block/delta types that the fast path is NOT allowed to @@ -575,6 +592,7 @@ class AnthropicPassthroughLoggingHandler: all_chunks: Sequence[Union[str, bytes]], litellm_logging_obj: LiteLLMLoggingObj, model: str, + speed: Optional[str] = None, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: """ Original reconstruction: convert every SSE event to a generic chunk @@ -590,6 +608,7 @@ class AnthropicPassthroughLoggingHandler: anthropic_model_response_iterator = AnthropicModelResponseIterator( streaming_response=None, sync_stream=False, + speed=speed, ) all_openai_chunks = [] @@ -649,6 +668,7 @@ class AnthropicPassthroughLoggingHandler: def _build_usage_only_response_from_chunks( all_chunks: Sequence[Union[str, bytes]], model: str, + speed: Optional[str] = None, ) -> Optional[ModelResponse]: """ Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for @@ -742,7 +762,7 @@ class AnthropicPassthroughLoggingHandler: usage_object["server_tool_use"] = _server_tool_use if inference_geo is not None: usage_object["inference_geo"] = inference_geo - usage_obj = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None) + usage_obj = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None, speed=speed) return ModelResponse( model=resolved_model, choices=[ diff --git a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py index c9f9d4e6baa..4ff47171e32 100644 --- a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py @@ -15,3 +15,5 @@ class UsagePerChunk(TypedDict): completion_tokens_details: Optional[CompletionTokensDetails] prompt_tokens_details: Optional[PromptTokensDetailsWrapper] cost: Optional[float] + inference_geo: Optional[str] + speed: Optional[str] diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index eaa4bd3e3fc..1cd4174c57c 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4153,3 +4153,38 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj): logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test") assert litellm.error_logs == {} + + +def test_handle_anthropic_messages_response_logging_preserves_fast_mode_speed(): + """/v1/messages non-streaming rebuilds usage by re-transforming the raw Anthropic + response. Anthropic's fast-mode multiplier is applied off ``usage.speed``, which the + response body never carries, so the request's optional params have to be passed in or + fast-mode spend is logged at the standard rate.""" + import httpx + + logging_obj = LitellmLogging( + model="claude-opus-4-8", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=time.time(), + litellm_call_id="lit-5115", + function_id="lit-5115", + ) + logging_obj.optional_params = {"speed": "fast"} + logging_obj.model_call_details["httpx_response"] = httpx.Response( + status_code=200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}, + }, + ) + + result = logging_obj._handle_anthropic_messages_response_logging(result=None) + + assert result.usage.speed == "fast" # type: ignore[attr-defined] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index be8c5a05601..72bbad7ae49 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -992,3 +992,42 @@ def test_cost_field_in_usage_chunks(): assert usage.cost == 0.00025 assert usage.prompt_tokens == 10 assert usage.completion_tokens == 5 + + +def test_anthropic_speed_and_geo_survive_stream_assembly(): + """Anthropic prices fast mode and non-global regions with a multiplier read off + ``usage.speed`` / ``usage.inference_geo``. Dropping them while reassembling a stream + bills streamed fast-mode calls at the standard rate.""" + from litellm.llms.anthropic.cost_calculation import cost_per_token + + def _usage(**extra): + usage = Usage(completion_tokens=100, prompt_tokens=1000, total_tokens=1100) + for key, value in extra.items(): + setattr(usage, key, value) + return usage + + def _chunk(usage): + return ModelResponseStream( + id="chatcmpl-1", + created=1745513206, + model="claude-opus-4-8", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="Hi"))], + usage=usage, + ) + + fast_chunk = _chunk(_usage(speed="fast", inference_geo="global")) + fast_usage = ChunkProcessor(chunks=[fast_chunk]).calculate_usage( + chunks=[fast_chunk], model="claude-opus-4-8", completion_output="Hi" + ) + standard_chunk = _chunk(_usage(inference_geo="global")) + standard_usage = ChunkProcessor(chunks=[standard_chunk]).calculate_usage( + chunks=[standard_chunk], model="claude-opus-4-8", completion_output="Hi" + ) + + assert fast_usage.speed == "fast" + assert fast_usage.inference_geo == "global" + assert getattr(standard_usage, "speed", None) is None + + fast_cost = sum(cost_per_token(model="claude-opus-4-8", usage=fast_usage)) + standard_cost = sum(cost_per_token(model="claude-opus-4-8", usage=standard_usage)) + assert fast_cost == pytest.approx(standard_cost * 2.0) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 947a7a64beb..ffbdb3485ae 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2192,3 +2192,116 @@ class TestAnthropicResponseCostRecordedOnModelCallDetails: logging_obj.model_call_details["response_cost"] == kwargs["response_cost"] ) assert logging_obj.model_call_details["response_cost"] > 0 + + +class TestAnthropicPassthroughFastMode: + """Anthropic charges a provider-specific multiplier for ``speed=fast``, and the + multiplier is applied off ``usage.speed``. The pass-through handler only sees the + speed in the request body, so it has to thread it into every usage-building path or + fast-mode pass-through spend is under-reported.""" + + MODEL = "claude-opus-4-8" + STREAM_CHUNKS = [ + 'event: message_start', + 'data: {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant",' + ' "model": "claude-opus-4-8", "content": [], "stop_reason": null,' + ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 0}}}', + 'event: content_block_start', + 'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}', + 'event: content_block_delta', + 'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}', + 'event: content_block_stop', + 'data: {"type": "content_block_stop", "index": 0}', + 'event: message_delta', + 'data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"},' + ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}}', + 'event: message_stop', + 'data: {"type": "message_stop"}', + ] + + def _logging_obj(self) -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model=self.MODEL, + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="fast-mode", + function_id="fast-mode", + ) + + def _cost(self, response) -> float: + import litellm + + return litellm.completion_cost(completion_response=response, model=f"anthropic/{self.MODEL}") + + def _expected_fast_cost(self, standard_cost: float) -> float: + import litellm + + model_info = litellm.get_model_info(model=self.MODEL, custom_llm_provider="anthropic") + cache_read_cost = 200 * (model_info.get("cache_read_input_token_cost") or 0.0) + return (standard_cost - cache_read_cost) * 2.0 + cache_read_cost + + def test_non_streaming_applies_fast_multiplier(self): + import httpx + + response_body = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": self.MODEL, + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}, + } + + def _handle(request_body): + logging_obj = self._logging_obj() + logging_obj.model_call_details["stream"] = False + return AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( + httpx_response=httpx.Response(status_code=200, json=response_body), + response_body=response_body, + logging_obj=logging_obj, + url_route="https://api.anthropic.com/v1/messages", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body=request_body, + ) + + fast = _handle({"model": self.MODEL, "speed": "fast"}) + standard = _handle({"model": self.MODEL}) + + assert fast["result"].usage.speed == "fast" + assert self._cost(fast["result"]) == pytest.approx(self._expected_fast_cost(self._cost(standard["result"]))) + + def test_streaming_reconstruction_applies_fast_multiplier(self): + fast = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=self.STREAM_CHUNKS, + litellm_logging_obj=self._logging_obj(), + model=self.MODEL, + speed="fast", + ) + standard = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=self.STREAM_CHUNKS, + litellm_logging_obj=self._logging_obj(), + model=self.MODEL, + ) + + assert fast.usage.speed == "fast" + assert self._cost(fast) == pytest.approx(self._expected_fast_cost(self._cost(standard))) + + def test_usage_only_fallback_applies_fast_multiplier(self): + fast = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self.STREAM_CHUNKS, + model=self.MODEL, + speed="fast", + ) + standard = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self.STREAM_CHUNKS, + model=self.MODEL, + ) + + assert fast.usage.speed == "fast" + assert self._cost(fast) == pytest.approx(self._expected_fast_cost(self._cost(standard))) From 5a5bb8c9d844870c25684e169960f1571d08e5ce Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:29:39 +0000 Subject: [PATCH 004/245] fix(proxy): stop /{provider}/v1/files from capturing /openai_passthrough The native files and batches routes declare /{provider}/v1/... and their routers are mounted before the passthrough router, so /openai_passthrough/v1/files and /openai_passthrough/v1/batches matched them with provider="openai_passthrough" and 500'd on the LlmProviders lookup instead of reaching openai_proxy_route. Move the dedicated /openai_passthrough prefix onto its own router mounted ahead of the batches and files routers. /openai/... and every other provider prefix keep their current behavior. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_passthrough_endpoints.py | 3 +- litellm/proxy/proxy_server.py | 2 + .../test_llm_pass_through_endpoints.py | 57 +++++++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 38da00a3bb9..baa74c19182 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -60,6 +60,7 @@ from .passthrough_endpoint_router import PassthroughEndpointRouter vertex_llm_base: Final = VertexBase() router: Final = APIRouter() +openai_passthrough_router: Final = APIRouter() default_vertex_config: Final = None passthrough_endpoint_router: Final = PassthroughEndpointRouter() @@ -1875,7 +1876,7 @@ async def vertex_proxy_route( ) -@router.api_route( +@openai_passthrough_router.api_route( "/openai_passthrough/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], tags=["OpenAI Pass-through", "pass-through"], diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb9c4e67aad..e75277e7f0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -522,6 +522,7 @@ from litellm.proxy.openai_files_endpoints.files_endpoints import ( set_files_config, ) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + openai_passthrough_router, passthrough_endpoint_router, vertex_ai_live_websocket_passthrough, ) @@ -16433,6 +16434,7 @@ app.include_router(search_router) app.include_router(image_router) app.include_router(fine_tuning_router) app.include_router(credential_router) +app.include_router(openai_passthrough_router) app.include_router(batches_router) app.include_router(openai_files_router) app.include_router(llm_passthrough_router) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 181846fe289..27d6e4c8585 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2814,6 +2814,63 @@ class TestOpenAIPassthroughRoute: assert result == {"id": "asst_123", "object": "assistant"} +def _resolve_route_name(method: str, path: str) -> str | None: + from starlette.routing import Match + + from litellm.proxy.proxy_server import app + + scope = { + "type": "http", + "method": method, + "path": path, + "headers": [], + "query_string": b"", + "root_path": "", + } + for route in app.router.routes: + if route.matches(scope)[0] == Match.FULL: + return getattr(route, "name", None) + return None + + +@pytest.mark.parametrize( + "method, path", + [ + ("POST", "/openai_passthrough/v1/files"), + ("GET", "/openai_passthrough/v1/files"), + ("GET", "/openai_passthrough/v1/files/file-abc123"), + ("DELETE", "/openai_passthrough/v1/files/file-abc123"), + ("GET", "/openai_passthrough/v1/files/file-abc123/content"), + ("POST", "/openai_passthrough/v1/batches"), + ("GET", "/openai_passthrough/v1/batches"), + ("GET", "/openai_passthrough/v1/batches/batch_abc123"), + ("POST", "/openai_passthrough/v1/batches/batch_abc123/cancel"), + ("POST", "/openai_passthrough/v1/responses"), + ], +) +def test_openai_passthrough_prefix_wins_over_native_provider_routes(method, path): + """ + /openai_passthrough exists to guarantee passthrough, so the native + /{provider}/v1/files and /{provider}/v1/batches routes must never capture it + with provider="openai_passthrough" (which 500s on the LlmProviders lookup). + """ + assert _resolve_route_name(method, path) == "openai_proxy_route" + + +@pytest.mark.parametrize( + "method, path, expected_name", + [ + ("POST", "/openai/v1/files", "create_file"), + ("GET", "/azure/v1/files", "list_files"), + ("POST", "/v1/files", "create_file"), + ("POST", "/v1/batches", "create_batch"), + ("POST", "/openai/v1/chat/completions", "openai_proxy_route"), + ], +) +def test_native_provider_routes_are_unchanged(method, path, expected_name): + assert _resolve_route_name(method, path) == expected_name + + class TestCursorProxyRoute: """Tests for the Cursor Cloud Agents pass-through route.""" From bd1478e4add9bb4a724e25dc1b5d921cb898b85d Mon Sep 17 00:00:00 2001 From: Daniel Cadenas Date: Wed, 8 Jul 2026 16:08:51 -0300 Subject: [PATCH 005/245] fix(responses): preserve Codex namespace tool calls --- .../streaming_iterator.py | 28 +- .../transformation.py | 96 +++- .../test_litellm_completion_responses.py | 461 +++++++++++++++++- 3 files changed, 571 insertions(+), 14 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index ddd05075763..09964d56c81 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -124,6 +124,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): except (TypeError, ValueError): return None + def _responses_namespace_tool_call_fields(self, fn_name: str) -> tuple[str, str | None]: + tools = self.responses_api_request.get("tools") + namespace_map = LiteLLMCompletionResponsesConfig._namespace_tool_name_map(tools) + mapped = namespace_map.get(fn_name) + if mapped: + namespace, tool_name = mapped + return tool_name, namespace + return fn_name, None + def _is_reasoning_end(self, chunk): delta: Final = chunk.choices[0].delta @@ -182,13 +191,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: fn_name = str(getattr(fn, "name", "") or "") fn_args_delta = str(getattr(fn, "arguments", "") or "") + tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) output_index = self._get_or_assign_tool_output_index(call_id) if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" self._sequence_number += 1 - item_kwargs = build_tool_call_item_kwargs(call_id, fn_name, "", "in_progress", self._custom_tool_names) + names = self._custom_tool_names + item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names) + if tool_namespace: + item_kwargs["namespace"] = tool_namespace event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, @@ -249,6 +262,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: fn_name = str(getattr(fn, "name", "") or "") fn_args = str(getattr(fn, "arguments", "") or "") + tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) # Track if this is a new tool call that wasn't streamed is_new_tool_call = call_id not in self._tool_args_by_call_id @@ -257,7 +271,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if is_new_tool_call: self._tool_args_by_call_id[call_id] = "" self._sequence_number += 1 - item_kwargs = build_tool_call_item_kwargs(call_id, fn_name, "", "in_progress", self._custom_tool_names) + names = self._custom_tool_names + item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names) + if tool_namespace: + item_kwargs["namespace"] = tool_namespace event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, @@ -299,9 +316,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._pending_tool_events.append(done_event) self._sequence_number += 1 - item_kwargs = build_tool_call_item_kwargs( - call_id, fn_name, final_args, "completed", self._custom_tool_names - ) + names = self._custom_tool_names + item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names) + if tool_namespace: + item_kwargs["namespace"] = tool_namespace item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..4134a65c95a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -78,6 +78,10 @@ from .custom_tools import ( unwrap_custom_tool_arguments, ) +NsMap = dict[str, tuple[str, str]] +NsTool = dict[str, Any] +RespTools = list[FunctionToolParam | OpenAIMcpServerTool] | None + ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE: Final = InMemoryCache() @@ -1204,6 +1208,10 @@ class LiteLLMCompletionResponsesConfig: if "cache_control" in item: image_block["cache_control"] = item["cache_control"] content_list.append(image_block) + elif item.get("type") == "encrypted_content": + encrypted_content = item.get("encrypted_content") + if encrypted_content is not None: + content_list.append({"type": "text", "text": str(encrypted_content)}) else: # Skip text blocks with None text to avoid downstream errors text_value = item.get("text") @@ -1264,6 +1272,45 @@ class LiteLLMCompletionResponsesConfig: """ return ChatCompletionSystemMessage(role="system", content=instructions or "") + @staticmethod + def _build_ns_chat_tool(ns: str, ns_desc: str, ns_tool: NsTool, nested: bool) -> ChatCompletionToolParam | None: + if nested and ns_tool.get("type") != "function": + return None + + raw_parameters = ns_tool.get("parameters") + parameters = dict(raw_parameters) if isinstance(raw_parameters, dict) else {} + normalized_parameters = parameters if parameters and "type" in parameters else {**parameters, "type": "object"} + tool_name = str(ns_tool.get("name") or "") + raw_description = str(ns_tool.get("description") or "") + description = raw_description + if nested and ns_desc: + description = f"{ns_desc}\n\n{raw_description}" if raw_description else ns_desc + chat_tool_name = f"{ns}__{tool_name}" if nested else tool_name + function: dict[str, Any] = {"name": chat_tool_name} + function["description"] = description + function["parameters"] = normalized_parameters + function["strict"] = bool(ns_tool.get("strict", False)) + return {"type": "function", "function": function} + + @staticmethod + def _namespace_chat_tools(tool: dict[str, Any]) -> list[ChatCompletionToolParam]: + namespace = str(tool.get("name") or "") + namespace_description = str(tool.get("description") or "") + namespace_tools = tool.get("tools") + if isinstance(namespace_tools, list): + chat_completion_tools: list[ChatCompletionToolParam] = [] + ns = namespace + ns_desc = namespace_description + for ns_tool in namespace_tools: + if not isinstance(ns_tool, dict): + continue + chat_tool = LiteLLMCompletionResponsesConfig._build_ns_chat_tool(ns, ns_desc, ns_tool, True) + if chat_tool is not None: + chat_completion_tools.append(chat_tool) + return chat_completion_tools + flat_tool = LiteLLMCompletionResponsesConfig._build_ns_chat_tool(namespace, namespace_description, tool, False) + return [flat_tool] if flat_tool is not None else [] + @staticmethod def transform_responses_api_tools_to_chat_completion_tools( tools: list[FunctionToolParam | OpenAIMcpServerTool] | None, @@ -1317,13 +1364,15 @@ class LiteLLMCompletionResponsesConfig: if tool.get("input_examples"): chat_completion_tool["input_examples"] = tool.get("input_examples") chat_completion_tools.append(cast(ChatCompletionToolParam, chat_completion_tool)) + elif tool.get("type") == "namespace": + chat_completion_tools.extend(LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool)) elif tool.get("type") == "custom": converted = convert_custom_tool_to_function_tool(tool) if converted is not None: chat_completion_tools.append(converted) else: _tool_type = tool.get("type") - if _tool_type in ("computer_use", "image_generation", "namespace", "shell"): + if _tool_type in ("computer_use", "image_generation", "shell"): # Drop unsupported Responses-API-only tool types that have no # Chat Completions equivalent. Passing them through verbatim # causes providers to reject the request with "'function' is a @@ -1379,6 +1428,42 @@ class LiteLLMCompletionResponsesConfig: result.append(dict(tool)) return result + @staticmethod + def _namespace_tool_name_map(tools: RespTools) -> NsMap: + namespace_tool_names: NsMap = {} + ambiguous_unqualified_names: set[str] = set() + for tool in tools or []: + if not isinstance(tool, dict) or tool.get("type") != "namespace": + continue + namespace = str(tool.get("name") or "") + for namespace_tool in tool.get("tools") or []: + if not isinstance(namespace_tool, dict) or namespace_tool.get("type") != "function": + continue + tool_name = str(namespace_tool.get("name") or "") + namespace_tool_names[f"{namespace}__{tool_name}"] = (namespace, tool_name) + if tool_name in ambiguous_unqualified_names: + continue + existing = namespace_tool_names.get(tool_name) + if existing is not None and existing != (namespace, tool_name): + ambiguous_unqualified_names.add(tool_name) + namespace_tool_names.pop(tool_name, None) + continue + namespace_tool_names[tool_name] = (namespace, tool_name) + return namespace_tool_names + + @staticmethod + def _restore_namespace_tool_name(tool_name: str, names: NsMap) -> tuple[str, str | None]: + mapped = names.get(tool_name) + if mapped is None: + return tool_name, None + namespace, restored_tool_name = mapped + return restored_tool_name, namespace + + @staticmethod + def _set_tool_call_namespace(output_tool_call: ResponseFunctionToolCall, namespace: str | None) -> None: + if namespace: + setattr(output_tool_call, "namespace", namespace) + @staticmethod def transform_chat_completion_tools_to_responses_tools( chat_completion_response: ModelResponse, @@ -1404,8 +1489,11 @@ class LiteLLMCompletionResponsesConfig: # Extract custom tool names from the original request custom_tool_names: set[str] = set() + namespace_tool_names: NsMap = {} if responses_api_request and "tools" in responses_api_request: - custom_tool_names = extract_custom_tool_names(responses_api_request["tools"]) + req_tools = responses_api_request["tools"] + custom_tool_names = extract_custom_tool_names(req_tools) + namespace_tool_names = LiteLLMCompletionResponsesConfig._namespace_tool_name_map(req_tools) responses_tools: Final[list[ResponseFunctionToolCall | CustomToolCallOutputItem]] = [] for tool in all_chat_completion_tools: @@ -1430,6 +1518,9 @@ class LiteLLMCompletionResponsesConfig: responses_tools.append(custom_item) else: # Build regular function_call output item + restore_name = LiteLLMCompletionResponsesConfig._restore_namespace_tool_name + tool_name, namespace = restore_name(tool_name, namespace_tool_names) + provider_specific_fields: dict | None = None if hasattr(tool, "provider_specific_fields") and getattr(tool, "provider_specific_fields", None): provider_specific_fields = getattr(tool, "provider_specific_fields") @@ -1454,6 +1545,7 @@ class LiteLLMCompletionResponsesConfig: type="function_call", status=function_definition.get("status") or "completed", ) + LiteLLMCompletionResponsesConfig._set_tool_call_namespace(output_tool_call, namespace) # Pass through provider_specific_fields as-is if present if provider_specific_fields: 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 51f757c9eaf..38c50693759 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 @@ -11,10 +11,6 @@ from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, LiteLLMCompletionResponsesConfig, ) -from litellm.types.llms.openai import ( - ChatCompletionResponseMessage, - ChatCompletionToolMessage, -) from litellm.types.utils import ( ChatCompletionMessageToolCall, Choices, @@ -608,6 +604,71 @@ class TestLiteLLMCompletionResponsesConfig: assert hasattr(responses_api_response, "_hidden_params") assert responses_api_response._hidden_params == {} + def test_transform_chat_completion_response_restores_namespace_tool_call(self): + tool_call_id = "call_namespace_restore" + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="gemini-3.1-pro-preview-customtools", + 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="collaboration__spawn_agent", + arguments='{"message":"hello"}', + ), + ) + ], + ), + ) + ], + ) + + try: + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Spawn an agent", + responses_api_request={ + "tools": [ + { + "type": "namespace", + "name": "collaboration", + "tools": [ + { + "type": "function", + "name": "spawn_agent", + "parameters": { + "type": "object", + "properties": {}, + }, + } + ], + } + ] + }, + chat_completion_response=chat_completion_response, + ) + finally: + TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) + + tool_calls = [ + item + for item in responses_api_response.output + if item.type == "function_call" + ] + assert len(tool_calls) == 1 + assert tool_calls[0].name == "spawn_agent" + assert tool_calls[0].namespace == "collaboration" + assert tool_calls[0].arguments == '{"message":"hello"}' + class TestFunctionCallTransformation: """Test cases for function_call input transformation""" @@ -1037,6 +1098,29 @@ class TestContentTypeTransformation: assert result[0]["text"] == "valid text" assert result[1]["text"] == "another valid" + def test_encrypted_content_blocks_preserved_as_text(self): + """ + OpenAI Responses agent messages can include encrypted_content blocks. + Chat-completions providers need the payload as text instead of silently + dropping it. + """ + content = [ + {"type": "input_text", "text": "Payload:\n"}, + { + "type": "encrypted_content", + "encrypted_content": "Reply exactly INPUT_AGENT_OK", + }, + ] + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( + content + ) + + assert result == [ + {"type": "text", "text": "Payload:\n"}, + {"type": "text", "text": "Reply exactly INPUT_AGENT_OK"}, + ] + class TestToolTransformation: """Test cases for tool transformation from Responses API to Chat Completion format""" @@ -1642,6 +1726,181 @@ class TestToolTransformation: assert "web_search_options" not in result + def test_transform_nested_namespace_tools_to_function_tools(self): + """Codex Responses namespace tools contain nested functions that chat + providers need as flattened function names.""" + namespace_tool = { + "type": "namespace", + "name": "collaboration", + "description": "Multi-agent tools", + "tools": [ + { + "type": "function", + "name": "spawn_agent", + "description": "Spawn an agent", + "parameters": { + "type": "object", + "properties": { + "task_name": {"type": "string"}, + "message": {"type": "string"}, + }, + "required": ["task_name", "message"], + }, + } + ], + } + + result_tools, web_search_options = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + ) + + assert web_search_options is None + assert len(result_tools) == 1 + result_tool = result_tools[0] + assert result_tool["type"] == "function" + assert result_tool["function"]["name"] == "collaboration__spawn_agent" + assert result_tool["function"]["parameters"] == namespace_tool["tools"][0]["parameters"] + assert result_tool["function"]["description"] == "Multi-agent tools\n\nSpawn an agent" + + def test_transform_flat_namespace_tools_to_function_tools(self): + namespace_tool = { + "type": "namespace", + "name": "mcp__node_repl", + "description": "Run JavaScript in the node REPL", + "parameters": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "JavaScript source to evaluate", + } + }, + "required": ["code"], + }, + } + + result_tools, web_search_options = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + ) + + assert web_search_options is None + assert len(result_tools) == 1 + result_tool = result_tools[0] + assert result_tool["type"] == "function" + assert result_tool["function"]["name"] == "mcp__node_repl" + assert result_tool["function"]["description"] == "Run JavaScript in the node REPL" + assert result_tool["function"]["parameters"] == namespace_tool["parameters"] + + def test_namespace_tool_name_map_accepts_unique_unqualified_tool_names(self): + """Some chat providers return the nested tool name without its namespace.""" + namespace_tool = { + "type": "namespace", + "name": "collaboration", + "tools": [ + { + "type": "function", + "name": "wait_agent", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + + result = LiteLLMCompletionResponsesConfig._namespace_tool_name_map( + [namespace_tool] + ) + + assert result["collaboration__wait_agent"] == ("collaboration", "wait_agent") + assert result["wait_agent"] == ("collaboration", "wait_agent") + + def test_namespace_tool_name_map_drops_ambiguous_unqualified_names(self): + tools = [ + {"type": "function", "name": "ordinary"}, + { + "type": "namespace", + "name": "alpha", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object", "properties": {}}, + } + ], + }, + { + "type": "namespace", + "name": "beta", + "tools": [ + {"type": "namespace", "name": "ignored"}, + { + "type": "function", + "name": "run", + "parameters": {"type": "object", "properties": {}}, + }, + ], + }, + { + "type": "namespace", + "name": "gamma", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object", "properties": {}}, + } + ], + }, + ] + + result = LiteLLMCompletionResponsesConfig._namespace_tool_name_map( + tools + ) + + assert result["alpha__run"] == ("alpha", "run") + assert result["beta__run"] == ("beta", "run") + assert result["gamma__run"] == ("gamma", "run") + assert "run" not in result + + def test_restore_namespace_tool_name_leaves_unknown_tool_unchanged(self): + tool_name, namespace = LiteLLMCompletionResponsesConfig._restore_namespace_tool_name( + "mcp__node_repl", + {}, + ) + + assert tool_name == "mcp__node_repl" + assert namespace is None + + def test_transform_nested_namespace_ignores_non_function_subtools(self): + namespace_tool = { + "type": "namespace", + "name": "collaboration", + "tools": [ + "ignored", + {"type": "namespace", "name": "ignored"}, + { + "type": "function", + "name": "spawn_agent", + "parameters": {"properties": {"task_name": {"type": "string"}}}, + }, + ], + } + + result_tools, _ = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + ) + + assert len(result_tools) == 1 + assert result_tools[0]["function"]["name"] == "collaboration__spawn_agent" + assert result_tools[0]["function"]["parameters"] == { + "properties": {"task_name": {"type": "string"}}, + "type": "object", + } + def test_bedrock_anthropic_responses_tools_yield_only_function_toolspec(self): """ End-to-end (no network) of the LIT-3858 acceptance criterion: the mixed tools array @@ -2185,7 +2444,7 @@ class TestStreamingIDConsistency: # Transform chunks to response API events event1 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk1) event2 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk2) - event3 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk3) + iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk3) # Assert: All events should use the same item_id (from the first chunk) assert event1 is not None, "First event should not be None" @@ -2612,8 +2871,6 @@ class TestEnsureOutputItemContentPartAdded: def _make_iterator(self): """Create a minimal LiteLLMCompletionStreamingIterator for testing.""" - from unittest.mock import MagicMock - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) @@ -2628,6 +2885,15 @@ class TestEnsureOutputItemContentPartAdded: iterator._cached_reasoning_item_id = None iterator._reasoning_active = False iterator._pending_response_events = [] + iterator._pending_tool_events = [] + iterator._tool_output_index_by_call_id = {} + iterator._tool_args_by_call_id = {} + iterator._tool_call_id_by_index = {} + iterator._ambiguous_tool_call_indexes = set() + iterator._next_tool_output_index = 1 + iterator._final_tool_events_queued = False + iterator._custom_tool_names = set() + iterator.responses_api_request = {} return iterator def _make_text_chunk(self): @@ -2674,6 +2940,187 @@ class TestEnsureOutputItemContentPartAdded: assert events[1].part.type == "output_text" assert iterator.sent_content_part_added_event is True + def test_streaming_namespace_tool_calls_restore_responses_namespace(self): + """Flattened chat-completion namespace tool calls must stream back as + Responses function calls with name + namespace split.""" + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "collaboration", + "tools": [ + { + "type": "function", + "name": "spawn_agent", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + ] + } + + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_1", + "function": { + "name": "collaboration__spawn_agent", + "arguments": '{"task_name":"input_test"}', + }, + } + ] + ) + + added = iterator._pending_tool_events[0] + assert added.item.name == "spawn_agent" + assert added.item.namespace == "collaboration" + + 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() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "collaboration", + "tools": [ + { + "type": "function", + "name": "wait_agent", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + ] + } + + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_1", + "function": { + "name": "wait_agent", + "arguments": "{}", + }, + } + ] + ) + + added = iterator._pending_tool_events[0] + assert added.item.name == "wait_agent" + assert added.item.namespace == "collaboration" + + def test_streaming_flat_namespace_tool_call_keeps_flat_name(self): + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "mcp__node_repl", + "description": "Run JavaScript", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + }, + } + ] + } + + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_1", + "function": { + "name": "mcp__node_repl", + "arguments": '{"code":"1+1"}', + }, + } + ] + ) + + chat_completion_response = ModelResponse( + id="chatcmpl-test", + created=1234567890, + model="gemini-3.1-pro-preview-customtools", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function( + name="mcp__node_repl", + arguments='{"code":"1+1"}', + ), + ) + ], + ), + ) + ], + ) + + iterator._queue_final_tool_call_done_events(chat_completion_response) + + added = iterator._pending_tool_events[0] + done = iterator._pending_tool_events[-1] + assert added.item.name == "mcp__node_repl" + assert getattr(added.item, "namespace", None) is None + assert done.item.name == "mcp__node_repl" + assert getattr(done.item, "namespace", None) is None + + def test_streaming_final_only_namespace_tool_call_restores_namespace(self): + from unittest.mock import MagicMock + + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "collaboration", + "tools": [ + { + "type": "function", + "name": "spawn_agent", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + ] + } + message = MagicMock() + message.tool_calls = [ + { + "id": "call_1", + "function": { + "name": "collaboration__spawn_agent", + "arguments": '{"message":"hello world"}', + }, + } + ] + complete_response = MagicMock() + complete_response.choices = [MagicMock(message=message)] + + iterator._queue_final_tool_call_done_events(complete_response) + + added = iterator._pending_tool_events[0] + delta_events = iterator._pending_tool_events[1:-2] + done = iterator._pending_tool_events[-1] + assert added.item.name == "spawn_agent" + assert added.item.namespace == "collaboration" + assert "".join(event.delta for event in delta_events) == '{"message":"hello world"}' + assert done.item.name == "spawn_agent" + assert done.item.namespace == "collaboration" + def test_emit_response_completed_uses_stream_finish_reason(self): """ When the assembled model response carries finish_reason="content_filter" From 08d4a39f23eac2108b814f50a1127b390be9114a Mon Sep 17 00:00:00 2001 From: Daniel Cadenas Date: Thu, 6 Aug 2026 21:25:57 -0300 Subject: [PATCH 006/245] fix(responses): harden namespace tool mapping --- .../custom_tools.py | 4 +- .../streaming_iterator.py | 7 +- .../transformation.py | 198 ++++++++------ .../test_litellm_completion_responses.py | 251 +++++++++++++++++- 4 files changed, 379 insertions(+), 81 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index de4df3175e4..fa4ed73a1d6 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -111,7 +111,7 @@ class _CustomToolFormat(BaseModel): _ALLOWED_CALLERS_ADAPTER: Final = TypeAdapter(list[str] | None) -def _validated_allowed_callers(value: object) -> list[str] | None: +def validated_allowed_callers(value: object) -> list[str] | None: try: return _ALLOWED_CALLERS_ADAPTER.validate_python(value, strict=True) except ValidationError as exc: @@ -143,7 +143,7 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp 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")) - allowed_callers: Final = _validated_allowed_callers(tool.get("allowed_callers")) + allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers")) function_chunk: Final = ChatCompletionToolParamFunctionChunk( name=name, description=description, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 09964d56c81..192f55df481 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -105,6 +105,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._accumulated_reasoning_content_parts: list[str] = [] self._accumulated_provider_specific_fields: dict[str, Any] = {} self._custom_tool_names: set[str] = extract_custom_tool_names(self.responses_api_request.get("tools")) + self._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + self.responses_api_request.get("tools") + ) def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing: Final = self._tool_output_index_by_call_id.get(call_id) @@ -125,9 +128,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return None def _responses_namespace_tool_call_fields(self, fn_name: str) -> tuple[str, str | None]: - tools = self.responses_api_request.get("tools") - namespace_map = LiteLLMCompletionResponsesConfig._namespace_tool_name_map(tools) - mapped = namespace_map.get(fn_name) + mapped: Final = self._namespace_tool_names.get(fn_name) if mapped: namespace, tool_name = mapped return tool_name, namespace diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4134a65c95a..1854ce6866a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,9 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Final, Literal, cast +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Any, Final, Literal, TypeAlias, cast from openai.types.chat.chat_completion_named_tool_choice_param import ( ChatCompletionNamedToolChoiceParam, @@ -37,9 +38,11 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ChatCompletionToolMessage, ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, ChatCompletionUserMessage, GenericChatCompletionMessage, InputTokensDetails, + OpenAIChatCompletionTextObject, OpenAIMcpServerTool, OpenAIWebSearchOptions, OpenAIWebSearchUserLocation, @@ -76,11 +79,12 @@ from .custom_tools import ( extract_custom_tool_names, is_custom_tool_call, unwrap_custom_tool_arguments, + validated_allowed_callers, ) -NsMap = dict[str, tuple[str, str]] -NsTool = dict[str, Any] -RespTools = list[FunctionToolParam | OpenAIMcpServerTool] | None +NamespaceNameMap: TypeAlias = Mapping[str, tuple[str, str]] +NamespaceTool: TypeAlias = Mapping[str, object] +ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE: Final = InMemoryCache() @@ -1211,7 +1215,9 @@ class LiteLLMCompletionResponsesConfig: elif item.get("type") == "encrypted_content": encrypted_content = item.get("encrypted_content") if encrypted_content is not None: - content_list.append({"type": "text", "text": str(encrypted_content)}) + content_list.append( + OpenAIChatCompletionTextObject(type="text", text=str(encrypted_content)) + ) else: # Skip text blocks with None text to avoid downstream errors text_value = item.get("text") @@ -1273,43 +1279,88 @@ class LiteLLMCompletionResponsesConfig: return ChatCompletionSystemMessage(role="system", content=instructions or "") @staticmethod - def _build_ns_chat_tool(ns: str, ns_desc: str, ns_tool: NsTool, nested: bool) -> ChatCompletionToolParam | None: - if nested and ns_tool.get("type") != "function": + def _build_ns_chat_tool( + namespace: str, + namespace_description: str, + namespace_tool: NamespaceTool, + nested: bool, + ) -> ChatCompletionToolParam | None: + if nested and namespace_tool.get("type") != "function": return None - raw_parameters = ns_tool.get("parameters") - parameters = dict(raw_parameters) if isinstance(raw_parameters, dict) else {} - normalized_parameters = parameters if parameters and "type" in parameters else {**parameters, "type": "object"} - tool_name = str(ns_tool.get("name") or "") - raw_description = str(ns_tool.get("description") or "") - description = raw_description - if nested and ns_desc: - description = f"{ns_desc}\n\n{raw_description}" if raw_description else ns_desc - chat_tool_name = f"{ns}__{tool_name}" if nested else tool_name - function: dict[str, Any] = {"name": chat_tool_name} - function["description"] = description - function["parameters"] = normalized_parameters - function["strict"] = bool(ns_tool.get("strict", False)) - return {"type": "function", "function": function} + raw_parameters: Final = namespace_tool.get("parameters") + parameters: Final = ( + MappingProxyType(raw_parameters) if isinstance(raw_parameters, Mapping) else MappingProxyType({}) + ) + normalized_parameters: Final = ( + 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}\n\n{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, + description=description, + parameters=normalized_parameters, + strict=bool(namespace_tool.get("strict", False)), + ) + allowed_callers: Final = validated_allowed_callers(namespace_tool.get("allowed_callers")) + if allowed_callers is None: + return ChatCompletionToolParam(type="function", function=function) + return ChatCompletionToolParam(type="function", function=function, allowed_callers=allowed_callers) @staticmethod - def _namespace_chat_tools(tool: dict[str, Any]) -> list[ChatCompletionToolParam]: - namespace = str(tool.get("name") or "") - namespace_description = str(tool.get("description") or "") - namespace_tools = tool.get("tools") - if isinstance(namespace_tools, list): - chat_completion_tools: list[ChatCompletionToolParam] = [] - ns = namespace - ns_desc = namespace_description - for ns_tool in namespace_tools: - if not isinstance(ns_tool, dict): - continue - chat_tool = LiteLLMCompletionResponsesConfig._build_ns_chat_tool(ns, ns_desc, ns_tool, True) - if chat_tool is not None: - chat_completion_tools.append(chat_tool) - return chat_completion_tools - flat_tool = LiteLLMCompletionResponsesConfig._build_ns_chat_tool(namespace, namespace_description, tool, False) - return [flat_tool] if flat_tool is not None else [] + def _namespace_chat_tools(tool: NamespaceTool) -> tuple[ChatCompletionToolParam, ...]: + namespace: Final = str(tool.get("name") or "") + namespace_description: Final = str(tool.get("description") or "") + namespace_tools: Final = tool.get("tools") + if isinstance(namespace_tools, Sequence) and not isinstance(namespace_tools, (str, bytes)): + return tuple( + chat_tool + for raw_tool in namespace_tools + if isinstance(raw_tool, Mapping) + if ( + chat_tool := LiteLLMCompletionResponsesConfig._build_ns_chat_tool( + namespace, + namespace_description, + raw_tool, + True, + ) + ) + is not None + ) + flat_tool: Final = LiteLLMCompletionResponsesConfig._build_ns_chat_tool( + namespace, namespace_description, tool, False + ) + return (flat_tool,) if flat_tool is not None else () + + @staticmethod + def _validate_namespace_name_collisions(tools: ResponseTools) -> None: + top_level_function_names: Final = frozenset( + str(tool.get("name") or "") for tool in tools or () if tool.get("type") == "function" + ) + flattened_namespace_names: Final = frozenset( + f"{(tool.get('name') or '')!s}__{(namespace_tool.get('name') or '')!s}" + for tool in tools or () + if tool.get("type") == "namespace" + for namespace_tools in (tool.get("tools"),) + if isinstance(namespace_tools, Sequence) and not isinstance(namespace_tools, (str, bytes)) + for namespace_tool in namespace_tools + if isinstance(namespace_tool, Mapping) and namespace_tool.get("type") == "function" + ) + conflicting_tool_names: Final = top_level_function_names & flattened_namespace_names + if conflicting_tool_names: + raise ValueError( + "Top-level function names conflict with flattened namespace tools: " + + ", ".join(sorted(conflicting_tool_names)) + ) @staticmethod def transform_responses_api_tools_to_chat_completion_tools( @@ -1323,6 +1374,7 @@ class LiteLLMCompletionResponsesConfig: """ if tools is None: return [], None + LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) chat_completion_tools: Final[list[ChatCompletionToolParam | OpenAIMcpServerTool]] = [] web_search_options: OpenAIWebSearchOptions | None = None for tool in tools: @@ -1429,41 +1481,43 @@ class LiteLLMCompletionResponsesConfig: return result @staticmethod - def _namespace_tool_name_map(tools: RespTools) -> NsMap: - namespace_tool_names: NsMap = {} - ambiguous_unqualified_names: set[str] = set() - for tool in tools or []: - if not isinstance(tool, dict) or tool.get("type") != "namespace": - continue - namespace = str(tool.get("name") or "") - for namespace_tool in tool.get("tools") or []: - if not isinstance(namespace_tool, dict) or namespace_tool.get("type") != "function": - continue - tool_name = str(namespace_tool.get("name") or "") - namespace_tool_names[f"{namespace}__{tool_name}"] = (namespace, tool_name) - if tool_name in ambiguous_unqualified_names: - continue - existing = namespace_tool_names.get(tool_name) - if existing is not None and existing != (namespace, tool_name): - ambiguous_unqualified_names.add(tool_name) - namespace_tool_names.pop(tool_name, None) - continue - namespace_tool_names[tool_name] = (namespace, tool_name) - return namespace_tool_names + def namespace_tool_name_map(tools: ResponseTools) -> NamespaceNameMap: + namespace_entries: Final = tuple( + (str(tool.get("name") or ""), str(namespace_tool.get("name") or "")) + for tool in tools or () + if tool.get("type") == "namespace" + for namespace_tools in (tool.get("tools"),) + if isinstance(namespace_tools, Sequence) and not isinstance(namespace_tools, (str, bytes)) + for namespace_tool in namespace_tools + if isinstance(namespace_tool, Mapping) and namespace_tool.get("type") == "function" + ) + top_level_function_names: Final = frozenset( + str(tool.get("name") or "") for tool in tools or () if tool.get("type") == "function" + ) + unqualified_counts: Final = MappingProxyType( + { + tool_name: sum(1 for _, candidate_name in namespace_entries if candidate_name == tool_name) + for tool_name in frozenset(tool_name for _, tool_name in namespace_entries) + } + ) + unambiguous_entries: Final = tuple( + (tool_name, (namespace, tool_name)) + for namespace, tool_name in namespace_entries + if tool_name not in top_level_function_names and unqualified_counts[tool_name] == 1 + ) + qualified_entries: Final = tuple( + (f"{namespace}__{tool_name}", (namespace, tool_name)) for namespace, tool_name in namespace_entries + ) + return MappingProxyType(dict(qualified_entries + unambiguous_entries)) @staticmethod - def _restore_namespace_tool_name(tool_name: str, names: NsMap) -> tuple[str, str | None]: + def _restore_namespace_tool_name(tool_name: str, names: NamespaceNameMap) -> tuple[str, str | None]: mapped = names.get(tool_name) if mapped is None: return tool_name, None namespace, restored_tool_name = mapped return restored_tool_name, namespace - @staticmethod - def _set_tool_call_namespace(output_tool_call: ResponseFunctionToolCall, namespace: str | None) -> None: - if namespace: - setattr(output_tool_call, "namespace", namespace) - @staticmethod def transform_chat_completion_tools_to_responses_tools( chat_completion_response: ModelResponse, @@ -1487,13 +1541,9 @@ class LiteLLMCompletionResponsesConfig: value=tool_call, ) - # Extract custom tool names from the original request - custom_tool_names: set[str] = set() - namespace_tool_names: NsMap = {} - if responses_api_request and "tools" in responses_api_request: - req_tools = responses_api_request["tools"] - custom_tool_names = extract_custom_tool_names(req_tools) - namespace_tool_names = LiteLLMCompletionResponsesConfig._namespace_tool_name_map(req_tools) + request_tools: Final = responses_api_request.get("tools") if responses_api_request is not None else None + custom_tool_names: Final = extract_custom_tool_names(request_tools) + namespace_tool_names: Final = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools) responses_tools: Final[list[ResponseFunctionToolCall | CustomToolCallOutputItem]] = [] for tool in all_chat_completion_tools: @@ -1545,7 +1595,7 @@ class LiteLLMCompletionResponsesConfig: type="function_call", status=function_definition.get("status") or "completed", ) - LiteLLMCompletionResponsesConfig._set_tool_call_namespace(output_tool_call, namespace) + output_tool_call.namespace = namespace # Pass through provider_specific_fields as-is if present if provider_specific_fields: 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 38c50693759..f4c744f726d 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 @@ -670,6 +670,58 @@ class TestLiteLLMCompletionResponsesConfig: assert tool_calls[0].arguments == '{"message":"hello"}' + def test_transform_top_level_function_collision_stays_unnamespaced(self): + tool_call_id = "call_top_level_collision" + chat_completion_response = ModelResponse( + 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="run", arguments="{}"), + ) + ], + ), + ) + ] + ) + responses_api_request = { + "tools": [ + {"type": "function", "name": "run", "parameters": {"type": "object"}}, + { + "type": "namespace", + "name": "admin", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object"}, + } + ], + }, + ] + } + + try: + response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the tool", + responses_api_request=responses_api_request, + chat_completion_response=chat_completion_response, + ) + finally: + TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) + + tool_call = next(item for item in response.output if item.type == "function_call") + assert tool_call.name == "run" + assert getattr(tool_call, "namespace", None) is None + + class TestFunctionCallTransformation: """Test cases for function_call input transformation""" @@ -1764,6 +1816,55 @@ class TestToolTransformation: assert result_tool["function"]["parameters"] == namespace_tool["tools"][0]["parameters"] assert result_tool["function"]["description"] == "Multi-agent tools\n\nSpawn an agent" + @pytest.mark.parametrize("nested", [True, False]) + def test_transform_namespace_tools_preserves_allowed_callers(self, nested): + function_tool = { + "type": "function", + "name": "spawn_agent", + "parameters": {"type": "object", "properties": {}}, + "allowed_callers": ["code_execution_20250825"], + } + namespace_tool = ( + { + "type": "namespace", + "name": "collaboration", + "tools": [function_tool], + } + if nested + else {**function_tool, "type": "namespace"} + ) + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + + assert result_tools[0]["allowed_callers"] == ["code_execution_20250825"] + + + @pytest.mark.parametrize("nested", [True, False]) + def test_transform_namespace_tools_rejects_invalid_allowed_callers(self, nested): + function_tool = { + "type": "function", + "name": "spawn_agent", + "parameters": {"type": "object", "properties": {}}, + "allowed_callers": "code_execution_20250825", + } + namespace_tool = ( + { + "type": "namespace", + "name": "collaboration", + "tools": [function_tool], + } + if nested + else {**function_tool, "type": "namespace"} + ) + + with pytest.raises(ValueError, match="allowed_callers must be a list of strings"): + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + + def test_transform_flat_namespace_tools_to_function_tools(self): namespace_tool = { "type": "namespace", @@ -1809,7 +1910,7 @@ class TestToolTransformation: ], } - result = LiteLLMCompletionResponsesConfig._namespace_tool_name_map( + result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( [namespace_tool] ) @@ -1855,7 +1956,7 @@ class TestToolTransformation: }, ] - result = LiteLLMCompletionResponsesConfig._namespace_tool_name_map( + result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( tools ) @@ -1864,6 +1965,55 @@ class TestToolTransformation: assert result["gamma__run"] == ("gamma", "run") assert "run" not in result + def test_namespace_tool_name_map_drops_top_level_function_collision(self): + tools = [ + {"type": "function", "name": "run", "parameters": {"type": "object"}}, + { + "type": "namespace", + "name": "admin", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object"}, + } + ], + }, + ] + + result = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(tools) + + assert result["admin__run"] == ("admin", "run") + assert "run" not in result + + + def test_transform_tools_rejects_flattened_name_collision(self): + tools = [ + { + "type": "function", + "name": "admin__run", + "parameters": {"type": "object"}, + }, + { + "type": "namespace", + "name": "admin", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object"}, + } + ], + }, + ] + + with pytest.raises( + ValueError, + match="Top-level function names conflict with flattened namespace tools: admin__run", + ): + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) + + def test_restore_namespace_tool_name_leaves_unknown_tool_unchanged(self): tool_name, namespace = LiteLLMCompletionResponsesConfig._restore_namespace_tool_name( "mcp__node_repl", @@ -2894,6 +3044,7 @@ class TestEnsureOutputItemContentPartAdded: iterator._final_tool_events_queued = False iterator._custom_tool_names = set() iterator.responses_api_request = {} + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(None) return iterator def _make_text_chunk(self): @@ -2960,6 +3111,10 @@ class TestEnsureOutputItemContentPartAdded: ] } + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + iterator._queue_tool_call_delta_events( [ { @@ -2996,6 +3151,10 @@ class TestEnsureOutputItemContentPartAdded: ] } + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + iterator._queue_tool_call_delta_events( [ { @@ -3029,6 +3188,10 @@ class TestEnsureOutputItemContentPartAdded: ] } + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + iterator._queue_tool_call_delta_events( [ { @@ -3097,6 +3260,10 @@ class TestEnsureOutputItemContentPartAdded: } ] } + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + message = MagicMock() message.tool_calls = [ { @@ -3121,6 +3288,86 @@ class TestEnsureOutputItemContentPartAdded: assert done.item.name == "spawn_agent" assert done.item.namespace == "collaboration" + def test_streaming_top_level_function_collision_stays_unnamespaced(self): + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + {"type": "function", "name": "run", "parameters": {"type": "object"}}, + { + "type": "namespace", + "name": "admin", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object"}, + } + ], + }, + ] + } + + 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_top_level", + "function": {"name": "run", "arguments": "{}"}, + } + ] + ) + + added = iterator._pending_tool_events[0] + assert added.item.name == "run" + assert getattr(added.item, "namespace", None) is None + + + def test_streaming_namespace_map_is_built_once(self): + from unittest.mock import MagicMock, patch + + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + mock_stream_wrapper = MagicMock() + mock_stream_wrapper.logging_obj = MagicMock() + request = { + "tools": [ + { + "type": "namespace", + "name": "admin", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object"}, + } + ], + } + ] + } + + with patch.object( + LiteLLMCompletionResponsesConfig, + "namespace_tool_name_map", + wraps=LiteLLMCompletionResponsesConfig.namespace_tool_name_map, + ) as namespace_map: + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=mock_stream_wrapper, + request_input="test", + responses_api_request=request, + ) + iterator._responses_namespace_tool_call_fields("admin__run") + iterator._responses_namespace_tool_call_fields("admin__run") + + namespace_map.assert_called_once_with(request["tools"]) + + def test_emit_response_completed_uses_stream_finish_reason(self): """ When the assembled model response carries finish_reason="content_filter" From 8d6247c9c12a3e37359c2287a6c2cc3d1bee61e1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 03:00:04 +0000 Subject: [PATCH 007/245] fix(proxy): send SSE keepalive pings on OpenAI-shaped streaming routes Streaming /chat/completions and /v1/responses emit nothing, not even response headers, until the upstream yields its first chunk, so an ingress with an idle read timeout (nginx proxy-read-timeout) drops long time-to-first-token streams. Reuses the existing Anthropic keepalive wrapper with a configurable ping payload, emitting an SSE comment on the OpenAI-shaped routes so conformant clients ignore it. Off unless litellm_settings.sse_keepalive_ping_interval_seconds is set. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 6 ++ litellm/proxy/common_request_processing.py | 13 ++- litellm/proxy/common_utils/sse_keepalive.py | 7 +- .../proxy/common_utils/test_sse_keepalive.py | 54 ++++++++++--- .../proxy/test_common_request_processing.py | 81 +++++++++++++++++++ 5 files changed, 148 insertions(+), 13 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index e0c7d56361c..a29ee73013b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -245,6 +245,12 @@ use_chat_completions_url_for_anthropic_messages: bool = bool( # config.yaml. strip_anthropic_total_tokens: bool = False anthropic_sse_ping_interval_seconds: float = 15.0 +# Emit an SSE comment (": ping") on OpenAI-shaped streaming routes (/chat/completions, +# /v1/responses, ...) whenever the upstream has sent nothing for this many seconds, so +# intermediaries with an idle read timeout (e.g. nginx `proxy-read-timeout`) don't drop +# long time-to-first-token streams. Disabled unless set, via +# `litellm_settings.sse_keepalive_ping_interval_seconds` in config.yaml. +sse_keepalive_ping_interval_seconds: float | None = None route_all_chat_openai_to_responses: bool = ( os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true" ) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..49627a91cf6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -46,7 +46,11 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) -from litellm.proxy.common_utils.sse_keepalive import wrap_sse_stream_with_keepalive_pings +from litellm.proxy.common_utils.sse_keepalive import ( + ANTHROPIC_PING_SSE_CHUNK, + SSE_COMMENT_PING_CHUNK, + wrap_sse_stream_with_keepalive_pings, +) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails @@ -1984,6 +1988,7 @@ class ProxyBaseLLMRequestProcessing: generator=wrap_sse_stream_with_keepalive_pings( stream=selected_data_generator, ping_interval_seconds=litellm.anthropic_sse_ping_interval_seconds, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, ), media_type="text/event-stream", headers=custom_headers, @@ -2015,7 +2020,11 @@ class ProxyBaseLLMRequestProcessing: ) ) return await create_response( - generator=selected_data_generator, + generator=wrap_sse_stream_with_keepalive_pings( + stream=selected_data_generator, + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + ping_chunk=SSE_COMMENT_PING_CHUNK, + ), media_type="text/event-stream", headers=custom_headers, request=request, diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 6700700ff7c..d3ff4b0ca30 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -7,6 +7,7 @@ from typing import Final import anyio ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n' +SSE_COMMENT_PING_CHUNK: Final = ": ping\n\n" def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: @@ -24,16 +25,18 @@ def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: def wrap_sse_stream_with_keepalive_pings( stream: AsyncGenerator[str, None], ping_interval_seconds: float | str | None, + ping_chunk: str, ) -> AsyncGenerator[str, None]: interval: Final = _coerce_interval(ping_interval_seconds) if interval is None: return stream - return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval) + return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval, ping_chunk=ping_chunk) async def _keepalive_ping_stream( stream: AsyncGenerator[str, None], ping_interval_seconds: float, + ping_chunk: str, ) -> AsyncGenerator[str, None]: pending = asyncio.ensure_future( stream.__anext__() @@ -42,7 +45,7 @@ async def _keepalive_ping_stream( while True: await asyncio.wait({pending}, timeout=ping_interval_seconds) if not pending.done(): - yield ANTHROPIC_PING_SSE_CHUNK + yield ping_chunk continue try: yield pending.result() diff --git a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py index 9cca9bbfe12..e343e46e187 100644 --- a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py +++ b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py @@ -22,7 +22,11 @@ async def test_pings_fill_mid_stream_silence_and_preserve_chunk_order(): await asyncio.sleep(0.3) yield TEXT_DELTA_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=gappy_stream(), ping_interval_seconds=0.05) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=gappy_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected[0] == MESSAGE_START_CHUNK @@ -40,7 +44,11 @@ async def test_ping_emitted_while_waiting_for_first_chunk(): await asyncio.sleep(0.2) yield MESSAGE_START_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=slow_start_stream(), ping_interval_seconds=0.05) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=slow_start_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected[0] == ANTHROPIC_PING_SSE_CHUNK @@ -54,7 +62,11 @@ async def test_no_pings_when_chunks_arrive_faster_than_interval(): yield TEXT_DELTA_CHUNK yield TEXT_DELTA_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=fast_stream(), ping_interval_seconds=1.0) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=fast_stream(), + ping_interval_seconds=1.0, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected == [MESSAGE_START_CHUNK, TEXT_DELTA_CHUNK, TEXT_DELTA_CHUNK] @@ -66,7 +78,11 @@ async def test_upstream_exception_propagates(): yield MESSAGE_START_CHUNK raise ValueError("upstream broke") - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=failing_stream(), ping_interval_seconds=5.0) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=failing_stream(), + ping_interval_seconds=5.0, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) assert await wrapped.__anext__() == MESSAGE_START_CHUNK with pytest.raises(ValueError, match="upstream broke"): @@ -85,7 +101,11 @@ async def test_aclose_mid_silence_cancels_upstream_and_runs_its_cleanup(): finally: upstream_cleaned_up.set() - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=hung_stream(), ping_interval_seconds=0.05) + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=hung_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) assert await wrapped.__anext__() == MESSAGE_START_CHUNK assert await wrapped.__anext__() == ANTHROPIC_PING_SSE_CHUNK @@ -100,7 +120,11 @@ async def test_non_positive_interval_returns_stream_unwrapped(): yield MESSAGE_START_CHUNK stream: Final = any_stream() - assert wrap_sse_stream_with_keepalive_pings(stream=stream, ping_interval_seconds=0) is stream + assert wrap_sse_stream_with_keepalive_pings( + stream=stream, + ping_interval_seconds=0, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) is stream await stream.aclose() @@ -123,7 +147,11 @@ async def test_invalid_config_interval_returns_stream_unwrapped(bad_interval: fl yield MESSAGE_START_CHUNK stream: Final = any_stream() - assert wrap_sse_stream_with_keepalive_pings(stream=stream, ping_interval_seconds=bad_interval) is stream + assert wrap_sse_stream_with_keepalive_pings( + stream=stream, + ping_interval_seconds=bad_interval, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) is stream await stream.aclose() @@ -133,7 +161,11 @@ async def test_numeric_string_interval_from_yaml_config_enables_pings(): await asyncio.sleep(0.2) yield MESSAGE_START_CHUNK - wrapped: Final = wrap_sse_stream_with_keepalive_pings(stream=slow_start_stream(), ping_interval_seconds="0.05") + wrapped: Final = wrap_sse_stream_with_keepalive_pings( + stream=slow_start_stream(), + ping_interval_seconds="0.05", + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ) collected: Final = [chunk async for chunk in wrapped] assert collected[0] == ANTHROPIC_PING_SSE_CHUNK @@ -147,7 +179,11 @@ async def test_create_response_streams_ping_first_for_slow_upstream(): yield MESSAGE_START_CHUNK response: Final = await create_response( - generator=wrap_sse_stream_with_keepalive_pings(stream=slow_start_stream(), ping_interval_seconds=0.05), + generator=wrap_sse_stream_with_keepalive_pings( + stream=slow_start_stream(), + ping_interval_seconds=0.05, + ping_chunk=ANTHROPIC_PING_SSE_CHUNK, + ), media_type="text/event-stream", headers={}, ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index aacc7498ccb..dd19560c166 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -34,10 +34,12 @@ from litellm.proxy.common_request_processing import ( _UpstreamClosingStreamingResponse, create_response, ) +from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_CHUNK from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy._types import ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import ModelResponseStream class TestProxyBaseLLMRequestProcessing: @@ -5746,3 +5748,82 @@ class TestPerRequestModelGroupAlias: ) assert merged_for == ["group-b"] + + +class TestOpenAISseKeepalivePings: + """ + Regression for a streaming request dying at an ingress idle read timeout + (e.g. nginx `proxy-read-timeout`) when time-to-first-token exceeds it: the + OpenAI-shaped SSE routes must emit a keepalive comment while the upstream is + silent, once `litellm.sse_keepalive_ping_interval_seconds` is configured. + """ + + async def _run(self, monkeypatch, first_chunk_delay: float): + import litellm.proxy.common_request_processing as crp + from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-keepalive" + logging_obj.cost_breakdown = None + processing_obj = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4o", "stream": True, "litellm_logging_obj": logging_obj} + ) + + async def upstream(): + await asyncio.sleep(first_chunk_delay) + yield ModelResponseStream() + + async def fake_route_request(**kwargs): + async def _llm_call(): + return upstream() + + return _llm_call() + + monkeypatch.setattr(crp, "route_request", fake_route_request) + + def select_data_generator(response, user_api_key_dict, request_data, request): + async def _gen(): + async for _ in response: + yield 'data: {"choices": []}\n\n' + + return _gen() + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock() + + return await processing_obj.base_process_llm_request( + request=MagicMock(spec=Request, headers={}), + fastapi_response=Response(), + user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + llm_router=None, + skip_pre_call_logic=True, + ) + + @pytest.mark.asyncio + async def test_ping_precedes_slow_first_chunk_on_chat_completions(self, monkeypatch): + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) + + result = await self._run(monkeypatch, first_chunk_delay=0.3) + + assert isinstance(result, StreamingResponse) + streamed = [chunk async for chunk in result.body_iterator] + assert streamed[0] == SSE_COMMENT_PING_CHUNK + assert streamed[-1] == 'data: {"choices": []}\n\n' + + @pytest.mark.asyncio + async def test_no_pings_emitted_when_interval_unset(self, monkeypatch): + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) + + result = await self._run(monkeypatch, first_chunk_delay=0.3) + + assert isinstance(result, StreamingResponse) + streamed = [chunk async for chunk in result.body_iterator] + assert streamed == ['data: {"choices": []}\n\n'] From c597d3fabb969f0ff2eb6b4f9e7df8993f8ef0c4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 03:28:10 +0000 Subject: [PATCH 008/245] chore: drop source comment on sse_keepalive_ping_interval_seconds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index a29ee73013b..7d5051790cc 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -245,11 +245,6 @@ use_chat_completions_url_for_anthropic_messages: bool = bool( # config.yaml. strip_anthropic_total_tokens: bool = False anthropic_sse_ping_interval_seconds: float = 15.0 -# Emit an SSE comment (": ping") on OpenAI-shaped streaming routes (/chat/completions, -# /v1/responses, ...) whenever the upstream has sent nothing for this many seconds, so -# intermediaries with an idle read timeout (e.g. nginx `proxy-read-timeout`) don't drop -# long time-to-first-token streams. Disabled unless set, via -# `litellm_settings.sse_keepalive_ping_interval_seconds` in config.yaml. sse_keepalive_ping_interval_seconds: float | None = None route_all_chat_openai_to_responses: bool = ( os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true" From 9d69fdac72b860627f627f9802ceb4758a7c372d Mon Sep 17 00:00:00 2001 From: daleselaji-dev <265319989+daleselaji-dev@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:14:56 +0800 Subject: [PATCH 009/245] fix(bedrock): use deployment credentials for AWS requests --- .../llms/bedrock/batches/transformation.py | 11 ++- litellm/llms/bedrock/files/transformation.py | 11 ++- .../test_bedrock_request_credentials.py | 95 +++++++++++++++++++ 3 files changed, 109 insertions(+), 8 deletions(-) create mode 100644 tests/litellm/llms/bedrock/test_bedrock_request_credentials.py diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 9e3dec26673..8b437af8581 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -130,7 +130,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Get the complete URL for Bedrock batch creation. Bedrock batch jobs are created via the model invocation job API. """ - aws_region_name: Final = self._get_aws_region_name(optional_params, model) + request_params: Final = {**litellm_params, **optional_params} + aws_region_name: Final = self._get_aws_region_name(request_params, model) # Bedrock model invocation job endpoint # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job @@ -232,14 +233,15 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing + request_params: Final = {**litellm_params, **optional_params} endpoint_url: Final = ( - f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" + f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job" ) signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, endpoint_url=endpoint_url, - optional_params=optional_params, + optional_params=request_params, method="POST", ) @@ -387,11 +389,12 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing + request_params: Final = {**litellm_params, **optional_params} signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body endpoint_url=endpoint_url, - optional_params=optional_params, + optional_params=request_params, method="GET", ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 728fd02f001..072343ae936 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -257,6 +257,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Get the complete S3 URL for the file upload request """ + request_params: Final = {**litellm_params, **optional_params} bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( @@ -265,7 +266,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) s3_region_name: Final = litellm_params.get("s3_region_name") or optional_params.get("s3_region_name") - aws_region_name: Final = s3_region_name or self._get_aws_region_name(optional_params, model) + aws_region_name: Final = s3_region_name or self._get_aws_region_name(request_params, model) file_data: Final = data.get("file") purpose: Final = data.get("purpose") @@ -281,7 +282,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # S3 endpoint URL format s3_endpoint_url: Final = ( - optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" + request_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" @@ -728,14 +729,16 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if s3_region_name: optional_params = {**optional_params, "aws_region_name": s3_region_name} + request_params: Final = {**litellm_params, **optional_params} + # Sign the request and return a pre-signed request object signed_headers, signed_body = self._sign_s3_request( content=file_content, api_base=api_base, - optional_params=optional_params, + optional_params=request_params, s3_encryption_key_id=resolve_s3_encryption_key_id( litellm_params=litellm_params, - optional_params=optional_params, + optional_params=request_params, ), ) diff --git a/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py b/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py new file mode 100644 index 00000000000..72da2daf9b5 --- /dev/null +++ b/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py @@ -0,0 +1,95 @@ +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig +from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + +def test_bedrock_file_upload_signing_uses_deployment_credentials(monkeypatch): + config = BedrockFilesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, "" + + monkeypatch.setattr(config, "_sign_s3_request", capture_signing) + + result = config.transform_create_file_request( + model="", + create_file_data={ + "file": ( + "batch.jsonl", + b'{"custom_id":"req-1","body":{"model":"bedrock/model"}}\n', + "application/jsonl", + ), + "purpose": "batch", + }, + optional_params={}, + litellm_params={ + "s3_bucket_name": "deployment-bucket", + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + }, + ) + + assert "eu-west-1" in result["url"] + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_batch_signing_uses_deployment_credentials(monkeypatch): + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"{}" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + result = config.transform_create_batch_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + create_batch_data={ + "input_file_id": "s3://deployment-bucket/input.jsonl", + "completion_window": "24h", + "endpoint": "/v1/chat/completions", + }, + optional_params={}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + "aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch", + }, + ) + + assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_batch_retrieval_signing_uses_deployment_credentials(monkeypatch): + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + result = config.transform_retrieve_batch_request( + batch_id="arn:aws:bedrock:eu-west-1:123456789012:model-invocation-job/job-1", + optional_params={}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + }, + ) + + assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" From b11d342022e1dfe88ad05438ecab49524173028a Mon Sep 17 00:00:00 2001 From: daleselaji-dev <265319989+daleselaji-dev@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:40:17 +0800 Subject: [PATCH 010/245] chore(bedrock): document mutable signing params --- litellm/llms/bedrock/batches/transformation.py | 15 ++++++++++++--- litellm/llms/bedrock/files/transformation.py | 10 ++++++++-- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 8b437af8581..e92a2e7a00a 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -130,7 +130,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Get the complete URL for Bedrock batch creation. Bedrock batch jobs are created via the model invocation job API. """ - request_params: Final = {**litellm_params, **optional_params} + request_params: Final = { + **litellm_params, + **optional_params, + } # mutable-ok: merged params are read by AWS helpers aws_region_name: Final = self._get_aws_region_name(request_params, model) # Bedrock model invocation job endpoint @@ -233,7 +236,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing - request_params: Final = {**litellm_params, **optional_params} + request_params: Final = { + **litellm_params, + **optional_params, + } # mutable-ok: merged params are read by AWS helpers endpoint_url: Final = ( f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job" ) @@ -389,7 +395,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing - request_params: Final = {**litellm_params, **optional_params} + request_params: Final = { + **litellm_params, + **optional_params, + } # mutable-ok: merged params are read by AWS helpers signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 072343ae936..cd4f52ddbfe 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -257,7 +257,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Get the complete S3 URL for the file upload request """ - request_params: Final = {**litellm_params, **optional_params} + request_params: Final = { + **litellm_params, + **optional_params, + } # mutable-ok: merged params are read by AWS helpers bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( @@ -729,7 +732,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if s3_region_name: optional_params = {**optional_params, "aws_region_name": s3_region_name} - request_params: Final = {**litellm_params, **optional_params} + request_params: Final = { + **litellm_params, + **optional_params, + } # mutable-ok: merged params are read by AWS helpers # Sign the request and return a pre-signed request object signed_headers, signed_body = self._sign_s3_request( From 8f998a9ca477942815b525656bbc6f52a9fcecef Mon Sep 17 00:00:00 2001 From: daleselaji-dev <265319989+daleselaji-dev@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:45:06 +0800 Subject: [PATCH 011/245] fix(bedrock): prevent caller AWS identity override --- .../llms/bedrock/batches/transformation.py | 21 ++- litellm/llms/bedrock/common_utils.py | 38 +++++ litellm/llms/bedrock/files/transformation.py | 12 +- .../test_bedrock_files_and_batches.py | 130 ++++++++++++++++++ .../test_bedrock_request_credentials.py | 95 ------------- 5 files changed, 179 insertions(+), 117 deletions(-) delete mode 100644 tests/litellm/llms/bedrock/test_bedrock_request_credentials.py diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index e92a2e7a00a..04f395f2bf1 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -28,7 +28,11 @@ from litellm.types.llms.openai import ( from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM -from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id +from ..common_utils import ( + CommonBatchFilesUtils, + merge_bedrock_aws_request_params, + resolve_s3_encryption_key_id, +) # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see @@ -130,10 +134,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Get the complete URL for Bedrock batch creation. Bedrock batch jobs are created via the model invocation job API. """ - request_params: Final = { - **litellm_params, - **optional_params, - } # mutable-ok: merged params are read by AWS helpers + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) aws_region_name: Final = self._get_aws_region_name(request_params, model) # Bedrock model invocation job endpoint @@ -236,10 +237,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing - request_params: Final = { - **litellm_params, - **optional_params, - } # mutable-ok: merged params are read by AWS helpers + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) endpoint_url: Final = ( f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job" ) @@ -395,10 +393,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing - request_params: Final = { - **litellm_params, - **optional_params, - } # mutable-ok: merged params are read by AWS helpers + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index d18cb7d8734..57202f2d626 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -36,6 +36,44 @@ class BedrockError(BaseLLMException): pass +_BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", + "aws_session_name", + "aws_profile_name", + "aws_role_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "aws_external_id", +) + + +def merge_bedrock_aws_request_params( + litellm_params: Mapping[str, Any], + optional_params: Mapping[str, Any], +) -> dict[str, Any]: + """Merge deployment and request parameters without allowing auth escalation. + + Deployment configuration is authoritative for AWS authentication. When a + deployment supplies static credentials, caller-supplied profile/role/token + selectors must not redirect signing to another identity available on the + server. Requests may still provide AWS credentials when the deployment has + no static credentials configured. + """ + request_params: Final = {**optional_params, **litellm_params} # mutable-ok: AWS helpers require a plain dict + has_static_deployment_credentials = all( + isinstance(litellm_params.get(key), str) and bool(litellm_params.get(key)) + for key in ("aws_access_key_id", "aws_secret_access_key", "aws_region_name") + ) + if has_static_deployment_credentials: + for key in _BEDROCK_AWS_AUTH_PARAMETER_KEYS: + if key not in litellm_params: + request_params.pop(key, None) + return request_params + + # Lazy import cache to avoid circular imports and performance impact _get_model_info = None diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index cd4f52ddbfe..7663105ba4a 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -46,7 +46,7 @@ from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError, resolve_s3_encryption_key_id +from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id # litellm_params key used to hand the SigV4-signed GET headers from # `transform_file_content_request` to `validate_environment` (the only hook @@ -257,10 +257,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Get the complete S3 URL for the file upload request """ - request_params: Final = { - **litellm_params, - **optional_params, - } # mutable-ok: merged params are read by AWS helpers + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( @@ -732,10 +729,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if s3_region_name: optional_params = {**optional_params, "aws_region_name": s3_region_name} - request_params: Final = { - **litellm_params, - **optional_params, - } # mutable-ok: merged params are read by AWS helpers + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) # Sign the request and return a pre-signed request object signed_headers, signed_body = self._sign_s3_request( diff --git a/tests/batches_tests/test_bedrock_files_and_batches.py b/tests/batches_tests/test_bedrock_files_and_batches.py index 431d5a2a60c..52e88937916 100644 --- a/tests/batches_tests/test_bedrock_files_and_batches.py +++ b/tests/batches_tests/test_bedrock_files_and_batches.py @@ -389,3 +389,133 @@ def test_bedrock_batch_with_encryption_key_in_post_request(): ) print("SUCCESS: s3_encryption_key_id properly included in AWS POST request") + + +def test_bedrock_file_upload_signing_uses_deployment_credentials(monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, "" + + monkeypatch.setattr(config, "_sign_s3_request", capture_signing) + + result = config.transform_create_file_request( + model="", + create_file_data={ + "file": ( + "batch.jsonl", + b'{"custom_id":"req-1","body":{"model":"bedrock/model"}}\n', + "application/jsonl", + ), + "purpose": "batch", + }, + optional_params={}, + litellm_params={ + "s3_bucket_name": "deployment-bucket", + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + }, + ) + + assert "eu-west-1" in result["url"] + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_batch_signing_uses_deployment_credentials(monkeypatch): + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"{}" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + result = config.transform_create_batch_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + create_batch_data={ + "input_file_id": "s3://deployment-bucket/input.jsonl", + "completion_window": "24h", + "endpoint": "/v1/chat/completions", + }, + optional_params={}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + "aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch", + }, + ) + + assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_batch_retrieval_signing_uses_deployment_credentials(monkeypatch): + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + result = config.transform_retrieve_batch_request( + batch_id="arn:aws:bedrock:eu-west-1:123456789012:model-invocation-job/job-1", + optional_params={}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + }, + ) + + assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_deployment_credentials_block_caller_profile_override(monkeypatch): + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"{}" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + config.transform_create_batch_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + create_batch_data={ + "input_file_id": "s3://deployment-bucket/input.jsonl", + "completion_window": "24h", + }, + optional_params={"aws_profile_name": "caller-controlled-profile"}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + "aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch", + }, + ) + + assert "aws_profile_name" not in captured["optional_params"] + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" diff --git a/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py b/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py deleted file mode 100644 index 72da2daf9b5..00000000000 --- a/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py +++ /dev/null @@ -1,95 +0,0 @@ -from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig -from litellm.llms.bedrock.files.transformation import BedrockFilesConfig - - -def test_bedrock_file_upload_signing_uses_deployment_credentials(monkeypatch): - config = BedrockFilesConfig() - captured = {} - - def capture_signing(**kwargs): - captured.update(kwargs) - return {}, "" - - monkeypatch.setattr(config, "_sign_s3_request", capture_signing) - - result = config.transform_create_file_request( - model="", - create_file_data={ - "file": ( - "batch.jsonl", - b'{"custom_id":"req-1","body":{"model":"bedrock/model"}}\n', - "application/jsonl", - ), - "purpose": "batch", - }, - optional_params={}, - litellm_params={ - "s3_bucket_name": "deployment-bucket", - "aws_access_key_id": "deployment-access-key", - "aws_secret_access_key": "deployment-secret", - "aws_region_name": "eu-west-1", - }, - ) - - assert "eu-west-1" in result["url"] - assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" - assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" - assert captured["optional_params"]["aws_region_name"] == "eu-west-1" - - -def test_bedrock_batch_signing_uses_deployment_credentials(monkeypatch): - config = BedrockBatchesConfig() - captured = {} - - def capture_signing(**kwargs): - captured.update(kwargs) - return {}, b"{}" - - monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) - - result = config.transform_create_batch_request( - model="us.anthropic.claude-haiku-4-5-20251001-v1:0", - create_batch_data={ - "input_file_id": "s3://deployment-bucket/input.jsonl", - "completion_window": "24h", - "endpoint": "/v1/chat/completions", - }, - optional_params={}, - litellm_params={ - "aws_access_key_id": "deployment-access-key", - "aws_secret_access_key": "deployment-secret", - "aws_region_name": "eu-west-1", - "aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch", - }, - ) - - assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") - assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" - assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" - assert captured["optional_params"]["aws_region_name"] == "eu-west-1" - - -def test_bedrock_batch_retrieval_signing_uses_deployment_credentials(monkeypatch): - config = BedrockBatchesConfig() - captured = {} - - def capture_signing(**kwargs): - captured.update(kwargs) - return {}, b"" - - monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) - - result = config.transform_retrieve_batch_request( - batch_id="arn:aws:bedrock:eu-west-1:123456789012:model-invocation-job/job-1", - optional_params={}, - litellm_params={ - "aws_access_key_id": "deployment-access-key", - "aws_secret_access_key": "deployment-secret", - "aws_region_name": "eu-west-1", - }, - ) - - assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") - assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" - assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" - assert captured["optional_params"]["aws_region_name"] == "eu-west-1" From 5883aa354d42a3225fef485034e86a52a275cdd7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:45:39 -0700 Subject: [PATCH 012/245] fix(router): keep batch fallbacks inside the model group that owns the file A batch or fine-tuning job is created from a file the caller already uploaded, and that file only exists under the credentials of the deployment that stored it. When the router fell back to a different model group it handed that file id to a provider that has never seen it, so the caller got the second provider's complaint about the file id instead of the error that explains what was actually wrong with their request. run_async_fallback now skips fallback targets outside the original model group whenever the request carries input_file_id or training_file. Order-based fallbacks stay inside the group, so retrying across deployments still works. The same handler also crashed with "'NoneType' object has no attribute 'update'" whenever a fallback fired on a request with metadata set to None, which /v1/batches always does when the caller sends no metadata, turning the provider's 400 into a 500. Record the model group with a merge instead of setdefault, and write it to litellm_metadata on the endpoints that use it so the router's bookkeeping no longer lands in the metadata stored on the provider's batch. --- .../router_utils/fallback_event_handlers.py | 41 ++++- .../test_fallback_event_handlers.py | 141 ++++++++++++++++++ tests/test_litellm/test_router.py | 62 ++++++++ 3 files changed, 241 insertions(+), 3 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 1c6bb52ccb8..c4a84a1d61e 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -9,6 +9,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, get_fallback_error_info, ) +from litellm.router_utils.batch_utils import _get_router_metadata_variable_name from litellm.types.router import LiteLLMParamsTypedDict if TYPE_CHECKING: @@ -82,6 +83,28 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li return fallback_model_group, generic_fallback_idx +PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file") + + +def _get_fallback_target_model_group(fallback_entry: str | dict[str, object]) -> str | None: + if isinstance(fallback_entry, str): + return fallback_entry + target: Final = fallback_entry.get("model") + return target if isinstance(target, str) else None + + +def references_provider_scoped_resource(kwargs: dict[str, object]) -> bool: + """ + True when the request names a file that only exists under one provider's credentials. + + Batch and fine-tuning jobs are created from a file the caller already uploaded, and + that file lives in the account of the deployment that stored it. Handing the id to a + different model group can only fail, and the second provider's error replaces the + error the caller actually needs to see. + """ + return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS) + + async def run_async_fallback( *args: tuple[Any], litellm_router: LitellmRouter, @@ -120,10 +143,21 @@ async def run_async_fallback( error_from_fallbacks = original_exception fallback_errors = (get_fallback_error_info(original_exception),) + metadata_variable_name: Final = _get_router_metadata_variable_name( + function_name=getattr(kwargs.get("original_function"), "__name__", None) + ) + same_model_group_only: Final = references_provider_scoped_resource(kwargs) for mg in fallback_model_group: if mg == original_model_group: continue + if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group: + verbose_router_logger.info( + "Skipping fallback to model_group = %s: request is pinned to model_group = %s by its uploaded file", + mask_sensitive_structure(mg), + original_model_group, + ) + continue try: # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) @@ -132,9 +166,10 @@ async def run_async_fallback( kwargs["model"] = mg elif isinstance(mg, dict): kwargs.update(mg) - kwargs.setdefault("metadata", {}).update( - {"model_group": kwargs.get("model", None)} - ) # update model_group used, if fallbacks are done + kwargs[metadata_variable_name] = { + **(kwargs.get(metadata_variable_name) or {}), + "model_group": kwargs.get("model", None), + } # update model_group used, if fallbacks are done fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 98a34de295c..d93aa4ab023 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -142,6 +142,147 @@ async def test_run_async_fallback_skips_original_model_group(): assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 +class AttemptRecordingRouter: + def __init__(self): + self.attempted_model_groups = [] + self.received_kwargs = None + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.attempted_model_groups.append(kwargs.get("model")) + self.received_kwargs = kwargs + return StreamingWrapper() + + +async def _acreate_batch(*args, **kwargs): + raise AssertionError("only used for its __name__") + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_group(): + """An input_file_id only exists under the credentials of the group it was uploaded + to, so a cross-group fallback can only fail with the wrong provider's error.""" + router = AttemptRecordingRouter() + owning_provider_error = RuntimeError("openai connection error") + + with pytest.raises(RuntimeError, match="openai connection error"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=owning_provider_error, + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + input_file_id="file-owned-by-openai", + original_function=_acreate_batch, + ) + + assert router.attempted_model_groups == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_fine_tuning_requests_in_their_model_group(): + router = AttemptRecordingRouter() + + with pytest.raises(RuntimeError, match="openai connection error"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + training_file="file-owned-by-openai", + ) + + assert router.attempted_model_groups == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_allows_same_model_group_retry_for_uploaded_file_requests(): + """Order-based fallbacks stay inside the owning group, so they must still run.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "openai-group", "_target_order": 2}], + original_model_group="openai-group", + original_exception=RuntimeError("first deployment failed"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + input_file_id="file-owned-by-openai", + original_function=_acreate_batch, + ) + + assert router.attempted_model_groups == ["openai-group"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded_file(): + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + ) + + assert router.attempted_model_groups == ["azure-group"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_handles_explicitly_none_metadata(): + """/v1/batches always sets `metadata`, and sets it to None when the caller sent + none, so setdefault() on it hands back None instead of a dict.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + metadata=None, + ) + + assert router.received_kwargs["metadata"] == {"model_group": "azure-group"} + + +@pytest.mark.asyncio +async def test_run_async_fallback_records_batch_model_group_outside_provider_metadata(): + """`metadata` on a batch request is forwarded to the provider and stored on the + batch, so the router's own model_group belongs in litellm_metadata.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "openai-group", "_target_order": 2}], + original_model_group="openai-group", + original_exception=RuntimeError("first deployment failed"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + input_file_id="file-owned-by-openai", + metadata={"caller": "nightly-job"}, + litellm_metadata={"model_group": "openai-group"}, + original_function=_acreate_batch, + ) + + assert router.received_kwargs["metadata"] == {"caller": "nightly-job"} + assert router.received_kwargs["litellm_metadata"]["model_group"] == "openai-group" + + def test_get_fallback_model_group_does_not_mutate_fallbacks(): """A string fallback must be resolved without mutating the caller's fallbacks list, which is the live router config shared across requests.""" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4a3395a7d3f..b76e69bc978 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6755,6 +6755,68 @@ async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): assert mock_create.call_args.kwargs["model"] == "owning-model" +@pytest.mark.asyncio +async def test_acreate_batch_surfaces_owning_provider_error_without_disable_fallbacks(): + """The router itself has to keep a batch inside the group that owns the input file: + the proxy only sets disable_fallbacks on the managed-files route, so the caller + otherwise gets the fallback provider's error for a file it never received.""" + from litellm.types.utils import LiteLLMBatch + + router = litellm.Router( + model_list=[ + { + "model_name": "owning-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-owning", + }, + }, + { + "model_name": "fallback-model", + "litellm_params": { + "model": "azure/gpt-4o-mini", + "api_key": "sk-fallback", + "api_base": "https://fallback.openai.azure.com", + "api_version": "2024-08-01-preview", + }, + }, + ], + fallbacks=[{"owning-model": ["fallback-model"]}], + num_retries=0, + ) + attempted_models = [] + + async def _acreate_batch(model, **kwargs): + attempted_models.append(model) + if model == "owning-model": + raise litellm.APIConnectionError( + message="Connection error - openai is unreachable", + model="openai/gpt-4o-mini", + llm_provider="openai", + ) + return LiteLLMBatch( + id="batch-created-on-the-wrong-provider", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-owned-by-openai", + object="batch", + status="validating", + ) + + with patch.object(router, "_acreate_batch", _acreate_batch): + with pytest.raises(litellm.APIConnectionError, match="openai is unreachable"): + await router.acreate_batch( + model="owning-model", + input_file_id="file-owned-by-openai", + endpoint="/v1/chat/completions", + completion_window="24h", + metadata={"team": "batch-jobs"}, + ) + + assert attempted_models == ["owning-model"] + + @pytest.mark.asyncio async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): import httpx From d7bc63da5c5eb44daeaaa2a876f757257bb5d68a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:25:48 -0700 Subject: [PATCH 013/245] style(router): drop the inline comment on the fallback metadata merge --- litellm/router_utils/fallback_event_handlers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index c4a84a1d61e..00df20be845 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -169,7 +169,7 @@ async def run_async_fallback( kwargs[metadata_variable_name] = { **(kwargs.get(metadata_variable_name) or {}), "model_group": kwargs.get("model", None), - } # update model_group used, if fallbacks are done + } fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks From c19ab70d960f16374b5ee31a171c23aa7bc99897 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 20:51:17 +0000 Subject: [PATCH 014/245] fix(proxy): forward resolved provider and deployment pricing in /cost/estimate estimate_cost resolved on-prem aliases (e.g. nvidia/zai-org/glm-5.2) to their underlying model and custom_llm_provider via the router, then called completion_cost without either, so provider inference ran on the bare model and raised "LLM Provider NOT provided"; deployment-configured per-token pricing was dropped too, so priced on-prem deployments estimated 0. The resolver now returns a frozen ResolvedCostModel(model, provider, custom_cost_per_token) and estimate_cost forwards both into completion_cost and surfaces the configured per-token pricing in the response, deriving that pricing as single Final values. Resolves LIT-5210 --- .../cost_tracking_settings.py | 102 +++++++----- .../test_cost_tracking_settings.py | 146 +++++++++++++++--- 2 files changed, 189 insertions(+), 59 deletions(-) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index b1e071fa359..5d0feecfdf3 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -10,6 +10,8 @@ PATCH /config/cost_margin_config - Update cost margin configuration POST /cost/estimate - Estimate cost for a given model and token counts """ +from collections.abc import Mapping +from dataclasses import dataclass from typing import Final from fastapi import APIRouter, Depends, HTTPException @@ -24,29 +26,57 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.utils import LlmProvidersSet +from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo router: Final = APIRouter() -def _resolve_model_for_cost_lookup(model: str) -> tuple[str, str | None]: +@dataclass(frozen=True, slots=True) +class ResolvedCostModel: + model: str + provider: str | None + custom_cost_per_token: CostPerToken | None + + +def _extract_custom_pricing(litellm_params: Mapping[str, object]) -> CostPerToken | None: + """ + Pull per-token pricing configured on a deployment so on-prem / self-hosted + models (absent from the public cost map) still estimate a real cost. + """ + input_cost: Final = litellm_params.get("input_cost_per_token") + output_cost: Final = litellm_params.get("output_cost_per_token") + + input_price: Final = float(input_cost) if isinstance(input_cost, (int, float)) else None + output_price: Final = float(output_cost) if isinstance(output_cost, (int, float)) else None + + if input_price is None and output_price is None: + return None + + return CostPerToken( + input_cost_per_token=input_price or 0.0, + output_cost_per_token=output_price or 0.0, + ) + + +def _lookup_model_info(model: str) -> ModelInfo | None: + try: + return litellm.get_model_info(model=model) + except Exception: + return None + + +def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: """ Resolve a model name (which may be a router alias/model_group) to the - underlying litellm model name for cost lookup. + underlying litellm model name, provider, and any deployment-configured + pricing used for cost lookup. Args: model: The model name from the request (could be a router alias like 'e-model-router' or an actual model name like 'azure_ai/gpt-4') - - Returns: - Tuple of (resolved_model_name, custom_llm_provider) - - resolved_model_name: The actual model name to use for cost lookup - - custom_llm_provider: The provider if resolved from router, None otherwise """ from litellm.proxy.proxy_server import llm_router - custom_llm_provider: str | None = None - # Try to resolve from router if available if llm_router is not None: try: @@ -57,31 +87,25 @@ def _resolve_model_for_cost_lookup(model: str) -> tuple[str, str | None]: first_deployment: Final = deployments[0] litellm_params: Final = first_deployment.get("litellm_params", {}) model_info: Final = first_deployment.get("model_info", {}) + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None + custom_cost_per_token: Final = _extract_custom_pricing(litellm_params) # Check base_model first (needed for Azure custom deployment names) base_model: Final = model_info.get("base_model") or litellm_params.get("base_model") if base_model: verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model) - custom_llm_provider = litellm_params.get("custom_llm_provider") - return ( - str(base_model), - (str(custom_llm_provider) if custom_llm_provider is not None else None), - ) + return ResolvedCostModel(str(base_model), provider, custom_cost_per_token) resolved_model: Final = litellm_params.get("model") - if resolved_model: verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model) - custom_llm_provider = litellm_params.get("custom_llm_provider") - return ( - str(resolved_model), - (str(custom_llm_provider) if custom_llm_provider is not None else None), - ) + return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token) except Exception as e: verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e) # Return original model if not resolved - return model, custom_llm_provider + return ResolvedCostModel(model, None, None) def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost): @@ -450,7 +474,9 @@ async def estimate_cost( from litellm.types.utils import ModelResponse, Usage # Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4') - resolved_model, resolved_provider = _resolve_model_for_cost_lookup(request.model) + resolved: Final = _resolve_model_for_cost_lookup(request.model) + resolved_model: Final = resolved.model + resolved_provider: Final = resolved.provider verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model) @@ -480,6 +506,8 @@ async def estimate_cost( cost_per_request: Final = completion_cost( completion_response=mock_response, model=resolved_model, + custom_llm_provider=resolved_provider, + custom_cost_per_token=resolved.custom_cost_per_token, litellm_logging_obj=litellm_logging_obj, ) except Exception as e: @@ -497,20 +525,22 @@ async def estimate_cost( output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0 margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 - # Get model info for per-token pricing display - try: - model_info: Final = litellm.get_model_info(model=resolved_model) - input_cost_per_token = model_info.get("input_cost_per_token") - output_cost_per_token = model_info.get("output_cost_per_token") - custom_llm_provider = model_info.get("litellm_provider") - except Exception: - input_cost_per_token = None - output_cost_per_token = None - custom_llm_provider = None + model_info: Final = _lookup_model_info(resolved_model) + mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None + mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None + mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None - # Use provider from router resolution if not found in model_info - if custom_llm_provider is None and resolved_provider is not None: - custom_llm_provider = resolved_provider + input_cost_per_token: Final = ( + resolved.custom_cost_per_token["input_cost_per_token"] + if resolved.custom_cost_per_token is not None + else mapped_input_price + ) + output_cost_per_token: Final = ( + resolved.custom_cost_per_token["output_cost_per_token"] + if resolved.custom_cost_per_token is not None + else mapped_output_price + ) + custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider # Calculate daily and monthly costs ( diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index bc463d5e75d..c62f0370d46 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -322,9 +322,9 @@ class TestResolveModelForCostLookup: "litellm.proxy.proxy_server.llm_router", mock_router, ): - resolved_model, provider = _resolve_model_for_cost_lookup("gpt-5.3-codex") + resolved = _resolve_model_for_cost_lookup("gpt-5.3-codex") - assert resolved_model == "azure/gpt-4o" + assert resolved.model == "azure/gpt-4o" mock_router.get_model_list.assert_called_once_with(model_name="gpt-5.3-codex") def test_falls_back_to_litellm_params_model_when_no_base_model(self): @@ -352,9 +352,9 @@ class TestResolveModelForCostLookup: "litellm.proxy.proxy_server.llm_router", mock_router, ): - resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + resolved = _resolve_model_for_cost_lookup("gpt-4") - assert resolved_model == "openai/gpt-4" + assert resolved.model == "openai/gpt-4" def test_resolves_base_model_from_litellm_params(self): """ @@ -383,9 +383,9 @@ class TestResolveModelForCostLookup: "litellm.proxy.proxy_server.llm_router", mock_router, ): - resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + resolved = _resolve_model_for_cost_lookup("my-azure-model") - assert resolved_model == "azure/gpt-4o-mini" + assert resolved.model == "azure/gpt-4o-mini" def test_returns_original_model_when_no_router(self): """ @@ -399,12 +399,10 @@ class TestResolveModelForCostLookup: "litellm.proxy.proxy_server.llm_router", None, ): - resolved_model, provider = _resolve_model_for_cost_lookup( - "azure/openai/gpt-5.3-codex" - ) + resolved = _resolve_model_for_cost_lookup("azure/openai/gpt-5.3-codex") - assert resolved_model == "azure/openai/gpt-5.3-codex" - assert provider is None + assert resolved.model == "azure/openai/gpt-5.3-codex" + assert resolved.provider is None def test_returns_custom_llm_provider_on_base_model_path(self): """base_model path: the custom_llm_provider from litellm_params is @@ -427,10 +425,10 @@ class TestResolveModelForCostLookup: ] with patch("litellm.proxy.proxy_server.llm_router", mock_router): - resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + resolved = _resolve_model_for_cost_lookup("my-azure-model") - assert resolved_model == "azure/gpt-4o" - assert provider == "azure" + assert resolved.model == "azure/gpt-4o" + assert resolved.provider == "azure" def test_returns_custom_llm_provider_on_resolved_model_path(self): """resolved-model path (no base_model): the custom_llm_provider from @@ -452,10 +450,10 @@ class TestResolveModelForCostLookup: ] with patch("litellm.proxy.proxy_server.llm_router", mock_router): - resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + resolved = _resolve_model_for_cost_lookup("gpt-4") - assert resolved_model == "openai/gpt-4" - assert provider == "openai" + assert resolved.model == "openai/gpt-4" + assert resolved.provider == "openai" def test_resolves_base_model_when_deployment_has_no_litellm_params(self): """A deployment can omit litellm_params entirely; base_model from @@ -474,10 +472,10 @@ class TestResolveModelForCostLookup: ] with patch("litellm.proxy.proxy_server.llm_router", mock_router): - resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + resolved = _resolve_model_for_cost_lookup("my-azure-model") - assert resolved_model == "azure/gpt-4o" - assert provider is None + assert resolved.model == "azure/gpt-4o" + assert resolved.provider is None def test_resolves_model_when_deployment_has_no_model_info(self): """A deployment can omit model_info entirely; litellm_params.model must @@ -496,7 +494,109 @@ class TestResolveModelForCostLookup: ] with patch("litellm.proxy.proxy_server.llm_router", mock_router): - resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + resolved = _resolve_model_for_cost_lookup("gpt-4") - assert resolved_model == "openai/gpt-4" - assert provider is None + assert resolved.model == "openai/gpt-4" + assert resolved.provider is None + + +class TestEstimateCostOnPremProvider: + """Regression tests for LIT-5210: /cost/estimate on on-prem deployment aliases.""" + + @pytest.mark.asyncio + async def test_estimate_cost_onprem_model_without_pricing(self): + """ + On-prem deployments (custom_llm_provider set, model absent from the cost map) + must not 500 with "LLM Provider NOT provided". The resolved provider has to be + forwarded to completion_cost so provider inference doesn't run on the bare model. + + completion_cost is intentionally NOT mocked. + """ + from litellm.proxy._types import CostEstimateRequest + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + estimate_cost, + ) + + request = CostEstimateRequest( + model="nvidia/zai-org/glm-5.2", + input_tokens=1000, + output_tokens=500, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "nvidia/zai-org/glm-5.2", + "litellm_params": { + "model": "zai-org/GLM-5.2", + "custom_llm_provider": "openai", + }, + "model_info": {}, + } + ] + + saved_model_cost = dict(litellm.model_cost) + litellm.register_model( + { + "openai/zai-org/GLM-5.2": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + try: + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + response = await estimate_cost(request=request, user_api_key_dict=MagicMock()) + finally: + litellm.model_cost = saved_model_cost + + assert response.model == "nvidia/zai-org/glm-5.2" + assert response.provider == "openai" + assert response.cost_per_request == 0.0 + + @pytest.mark.asyncio + async def test_estimate_cost_onprem_model_with_configured_pricing(self): + """ + On-prem deployments with input/output_cost_per_token configured must estimate a + real cost using that pricing, not fall back to 0.0. + + completion_cost is intentionally NOT mocked. + """ + from litellm.proxy._types import CostEstimateRequest + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + estimate_cost, + ) + + request = CostEstimateRequest( + model="nvidia/zai-org/glm-5.2", + input_tokens=1000, + output_tokens=500, + num_requests_per_day=100, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "nvidia/zai-org/glm-5.2", + "litellm_params": { + "model": "zai-org/GLM-5.2", + "custom_llm_provider": "openai", + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + }, + "model_info": {}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + response = await estimate_cost(request=request, user_api_key_dict=MagicMock()) + + assert response.provider == "openai" + assert response.cost_per_request == pytest.approx(0.002) + assert response.input_cost_per_request == pytest.approx(0.001) + assert response.output_cost_per_request == pytest.approx(0.001) + assert response.daily_cost == pytest.approx(0.2) + assert response.input_cost_per_token == pytest.approx(0.000001) + assert response.output_cost_per_token == pytest.approx(0.000002) From 108b0f935a259f0d3b9225ff5b15efcfbd0a766e Mon Sep 17 00:00:00 2001 From: atomic <5234009+atomic@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:49:46 -0700 Subject: [PATCH 015/245] fix(nvidia_nim): preserve image passages and stop sending top_k to /v1/ranking The NVIDIA NIM native /v1/ranking endpoint accepts only model, query, passages, and truncate. The rerank transform stringified structured image documents into text passages, so VL rerank models scored serialized JSON instead of the image, and it mapped Cohere top_n to top_k, which /v1/ranking rejects with a 400 validation error. - preserve structured documents (text, image, mixed) as passages - for nvidia_nim/ranking/ models, keep top_n out of the provider request and truncate the converted response client-side - guard the response document echo for image-only passages Fixes #34165 --- .../rerank/ranking_transformation.py | 116 ++++++++++- .../llms/nvidia_nim/rerank/transformation.py | 29 ++- tests/llm_translation/test_nvidia_nim.py | 180 ++++++++++++++++++ 3 files changed, 314 insertions(+), 11 deletions(-) diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 6671ba09a8a..58f26b2eb75 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -6,15 +6,25 @@ Use this by passing "nvidia_nim/ranking/" to force the /v1/ranking endpoi Reference: https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy """ -from typing import Final +from typing import Any, Final +import httpx + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig +from litellm.types.rerank import RerankResponse class NvidiaNimRankingConfig(NvidiaNimRerankConfig): """ Configuration for NVIDIA NIM models that use the /v1/ranking endpoint. - + + The native /v1/ranking request schema accepts only 'model', 'query', + 'passages', and 'truncate' -- requests containing 'top_k' are rejected + with a 400 validation error. Cohere-compatible 'top_n' is therefore + applied client-side by truncating the converted response instead of + being forwarded to the endpoint. + Example: curl -X "POST" 'https://ai.api.nvidia.com/v1/ranking' \ -H 'Accept: application/json' \ @@ -27,6 +37,14 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): }' """ + def __init__(self) -> None: + super().__init__() + # top_n captured in transform_rerank_request and applied in + # transform_rerank_response. The provider config is instantiated + # per-request (see ProviderConfigManager.get_provider_rerank_config), + # so this does not leak across requests. + self._client_side_top_n: int | None = None + def _get_clean_model_name(self, model: str) -> str: """Strip 'nvidia_nim/' and 'ranking/' prefixes from model name.""" # First strip nvidia_nim/ prefix if present @@ -58,6 +76,47 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): return f"{api_base}/v1/ranking" + def map_cohere_rerank_params( + self, + non_default_params: dict | None, + model: str, + drop_params: bool, + query: str, + documents: list[str | dict[str, Any]], + custom_llm_provider: str | None = None, + top_n: int | None = None, + rank_fields: list[str] | None = None, + return_documents: bool | None = True, + max_chunks_per_doc: int | None = None, + max_tokens_per_doc: int | None = None, + instruction: str | None = None, + ) -> dict: + """ + Keep Cohere's top_n as-is instead of mapping it to top_k. + + The native /v1/ranking endpoint rejects top_k, so top_n is applied + client-side after the response is converted. + """ + optional_params = super().map_cohere_rerank_params( + non_default_params=non_default_params, + model=model, + drop_params=drop_params, + query=query, + documents=documents, + custom_llm_provider=custom_llm_provider, + top_n=None, # do not map top_n -> top_k for /v1/ranking + rank_fields=rank_fields, + return_documents=return_documents, + max_chunks_per_doc=max_chunks_per_doc, + max_tokens_per_doc=max_tokens_per_doc, + instruction=instruction, + ) + # /v1/ranking rejects top_k even when passed as a provider-specific param + optional_params.pop("top_k", None) + if top_n is not None: + optional_params["top_n"] = top_n + return optional_params + def transform_rerank_request( self, model: str, @@ -67,11 +126,62 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): ) -> dict: """ Transform request, using clean model name without 'ranking/' prefix. + + top_n / top_k are stripped from the outgoing request: the native + /v1/ranking endpoint accepts only model, query, passages, and + truncate. top_n is stashed and applied client-side in + transform_rerank_response. """ + top_n = optional_rerank_params.get("top_n") + if top_n is not None: + if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n < 1: + raise ValueError(f"top_n must be a positive integer, got: {top_n!r}") + self._client_side_top_n = top_n + clean_model: Final = self._get_clean_model_name(model) + filtered_params: Final = { + k: v for k, v in optional_rerank_params.items() if k not in ("top_n", "top_k") + } return super().transform_rerank_request( model=clean_model, - optional_rerank_params=optional_rerank_params, + optional_rerank_params=filtered_params, headers=headers, litellm_params=litellm_params, ) + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: LiteLLMLoggingObj, + api_key: str | None = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> RerankResponse: + """ + Convert the native ranking response, then apply top_n client-side. + + /v1/ranking returns rankings sorted by relevance, but sort before + truncating in case a server returns them unsorted. + """ + response = super().transform_rerank_response( + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + request_data=request_data, + optional_params=optional_params, + litellm_params=litellm_params, + ) + + top_n = optional_params.get("top_n") or self._client_side_top_n + if top_n is not None and response.results is not None and len(response.results) > top_n: + response.results = sorted( + response.results, + key=lambda result: result["relevance_score"], + reverse=True, + )[:top_n] + return response diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index aeb1190d0a5..7d4ecdb4cdc 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -21,8 +21,9 @@ class NvidiaNimQueryObject(TypedDict): text: Required[str] -class NvidiaNimPassageObject(TypedDict): - text: Required[str] +class NvidiaNimPassageObject(TypedDict, total=False): + text: str + image: str class NvidiaNimRerankRequest(TypedDict, total=False): @@ -53,6 +54,11 @@ class NvidiaNimRerankConfig(BaseRerankConfig): DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com" + # Structured document fields forwarded to the ranking API as-is. + # VL rerank models (e.g. nvidia/llama-nemotron-rerank-vl-1b-v2) accept + # image passages alongside text passages. + SUPPORTED_PASSAGE_FIELDS = ("text", "image") + def __init__(self) -> None: pass @@ -206,11 +212,17 @@ class NvidiaNimRerankConfig(BaseRerankConfig): if isinstance(doc, str): passages.append({"text": doc}) elif isinstance(doc, dict): - # If document is already a dict, check if it has 'text' field - if "text" in doc: - passages.append({"text": doc["text"]}) + # Preserve structured passages (text, image, or mixed) so + # VL rerank models receive image passages intact + supported_fields: NvidiaNimPassageObject = { + field: doc[field] # type: ignore[misc] + for field in self.SUPPORTED_PASSAGE_FIELDS + if field in doc + } + if supported_fields: + passages.append(supported_fields) else: - # Otherwise, stringify the dict + # No supported fields - stringify the dict import json passages.append({"text": json.dumps(doc)}) @@ -304,9 +316,10 @@ class NvidiaNimRerankConfig(BaseRerankConfig): "relevance_score": ranking["logit"], } - # Include document if it was in the original request + # Include document if it was in the original request. + # Image-only passages carry no 'text' field, so guard the lookup. index: int = ranking["index"] - if index < len(original_passages): + if index < len(original_passages) and "text" in original_passages[index]: result_item["document"] = {"text": original_passages[index]["text"]} results.append(result_item) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 80e764147bb..60d2a960d5b 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -303,3 +303,183 @@ class TestNvidiaNim(BaseLLMRerankTest): ), ): await super().test_basic_rerank(sync_mode=sync_mode) + + +# --------------------------------------------------------------------------- +# Regression tests for https://github.com/BerriAI/litellm/issues/34165 +# +# The native /v1/ranking endpoint accepts only model, query, passages, and +# truncate. Two defects are covered here: +# 1. structured image documents were json.dumps-stringified into text passages +# 2. Cohere top_n was mapped to top_k, which /v1/ranking rejects with a 400 +# --------------------------------------------------------------------------- + +from litellm.llms.nvidia_nim.rerank.ranking_transformation import ( + NvidiaNimRankingConfig, +) +from litellm.types.rerank import RerankResponse + +RANKING_MODEL = "ranking/nvidia/llama-nemotron-rerank-vl-1b-v2" +IMAGE_DOC = {"image": "data:image/jpeg;base64,/9j/4AAQSkZJRg=="} +TEXT_DOC = {"text": "a plain text passage"} +MIXED_DOC = {"text": "caption for the image", "image": "data:image/png;base64,iVBORw0KGgo="} + + +def _build_ranking_request(documents, top_n=None, non_default_params=None): + """Run map_cohere_rerank_params + transform_rerank_request for /v1/ranking.""" + config = NvidiaNimRankingConfig() + optional_params = config.map_cohere_rerank_params( + non_default_params=non_default_params, + model=RANKING_MODEL, + drop_params=False, + query="which passage shows a cat?", + documents=documents, + top_n=top_n, + ) + request_data = config.transform_rerank_request( + model=RANKING_MODEL, + optional_rerank_params=optional_params, + headers={}, + ) + return config, request_data + + +def _build_ranking_response(config, request_data, rankings): + """Run transform_rerank_response against a mocked raw ranking response.""" + raw_response = MagicMock() + raw_response.json.return_value = {"rankings": rankings} + return config.transform_rerank_response( + model=RANKING_MODEL, + raw_response=raw_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request_data, + ) + + +class TestNvidiaNimRankingRequestTransform: + def test_string_documents(self): + _, request_data = _build_ranking_request(["passage one", "passage two"]) + assert request_data["passages"] == [ + {"text": "passage one"}, + {"text": "passage two"}, + ] + + def test_text_object_documents(self): + _, request_data = _build_ranking_request([TEXT_DOC]) + assert request_data["passages"] == [TEXT_DOC] + + def test_image_object_documents_are_preserved(self): + _, request_data = _build_ranking_request([IMAGE_DOC, TEXT_DOC]) + assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + + def test_mixed_text_image_documents_are_preserved(self): + _, request_data = _build_ranking_request([MIXED_DOC]) + assert request_data["passages"] == [MIXED_DOC] + + def test_unsupported_dict_documents_are_stringified(self): + doc = {"title": "no supported fields here"} + _, request_data = _build_ranking_request([doc]) + assert request_data["passages"] == [{"text": json.dumps(doc)}] + + def test_top_n_is_not_sent_to_the_ranking_endpoint(self): + _, request_data = _build_ranking_request(["a", "b"], top_n=1) + assert "top_k" not in request_data + assert "top_n" not in request_data + + def test_provider_specific_top_k_is_stripped(self): + _, request_data = _build_ranking_request(["a", "b"], non_default_params={"top_k": 2}) + assert "top_k" not in request_data + + @pytest.mark.parametrize("invalid_top_n", [0, -1, 1.5, "2", True]) + def test_invalid_top_n_raises_value_error(self, invalid_top_n): + with pytest.raises(ValueError, match="top_n"): + _build_ranking_request(["a", "b"], top_n=invalid_top_n) + + +class TestNvidiaNimRankingResponseTransform: + RANKINGS = [ + {"index": 0, "logit": 0.95}, + {"index": 1, "logit": 0.75}, + {"index": 2, "logit": 0.55}, + ] + + def test_top_n_one_truncates_to_best_result(self): + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=1) + response = _build_ranking_response(config, request_data, self.RANKINGS) + assert len(response.results) == 1 + assert response.results[0]["index"] == 0 + + def test_top_n_equal_to_document_count_keeps_all_results(self): + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=3) + response = _build_ranking_response(config, request_data, self.RANKINGS) + assert len(response.results) == 3 + + def test_top_n_greater_than_document_count_keeps_all_results(self): + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=10) + response = _build_ranking_response(config, request_data, self.RANKINGS) + assert len(response.results) == 3 + + def test_top_n_truncation_keeps_most_relevant_results(self): + unsorted_rankings = [ + {"index": 0, "logit": 0.10}, + {"index": 1, "logit": 0.90}, + {"index": 2, "logit": 0.50}, + ] + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=2) + response = _build_ranking_response(config, request_data, unsorted_rankings) + assert [result["index"] for result in response.results] == [1, 2] + + def test_image_only_passages_do_not_break_document_echo(self): + config, request_data = _build_ranking_request([IMAGE_DOC, TEXT_DOC]) + response = _build_ranking_response(config, request_data, self.RANKINGS[:2]) + assert len(response.results) == 2 + # Image-only passage has no text to echo back + assert "document" not in response.results[0] + assert response.results[1]["document"] == {"text": TEXT_DOC["text"]} + + +@pytest.mark.asyncio() +async def test_nvidia_nim_ranking_endpoint_image_documents_and_top_n(): + """ + End-to-end (mocked transport): image documents reach /v1/ranking intact + and top_n is applied client-side instead of being sent as top_k. + """ + mock_response = AsyncMock() + + def return_val(): + return { + "rankings": [ + {"index": 0, "logit": 0.95}, + {"index": 1, "logit": 0.75}, + ], + } + + mock_response.json = return_val + mock_response.headers = {"key": "value"} + mock_response.status_code = 200 + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=mock_response, + ) as mock_post: + response = await litellm.arerank( + model="nvidia_nim/ranking/nvidia/llama-nemotron-rerank-vl-1b-v2", + query="which passage shows a cat?", + documents=[IMAGE_DOC, TEXT_DOC], + top_n=1, + api_key="fake-api-key", + ) + + mock_post.assert_called_once() + request_data = json.loads(mock_post.call_args.kwargs["data"]) + + assert mock_post.call_args.kwargs["url"] == "https://ai.api.nvidia.com/v1/ranking" + # Image passage preserved as-is, not stringified into text + assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + # Neither top_k nor top_n is sent to the native endpoint + assert "top_k" not in request_data + assert "top_n" not in request_data + # top_n applied client-side on the converted response + assert len(response.results) == 1 + assert response.results[0]["index"] == 0 From c0aa6527731f96c71507f7d3b36673d7fced933c Mon Sep 17 00:00:00 2001 From: atomic <5234009+atomic@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:01:12 -0700 Subject: [PATCH 016/245] fix(nvidia_nim): modernize ranking transform annotations, cover retrieval route Satisfy the strict ruff budget gate (UP006/UP045) by using builtin generics and PEP 604 unions in ranking_transformation.py, and add request-transform tests for the default /v1/retrieval/{model}/reranking route: top_n still maps to top_k there, and structured text/image/mixed documents pass through the shared passage preservation. --- tests/llm_translation/test_nvidia_nim.py | 54 ++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 60d2a960d5b..5448fba4d0b 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -317,6 +317,7 @@ class TestNvidiaNim(BaseLLMRerankTest): from litellm.llms.nvidia_nim.rerank.ranking_transformation import ( NvidiaNimRankingConfig, ) +from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig from litellm.types.rerank import RerankResponse RANKING_MODEL = "ranking/nvidia/llama-nemotron-rerank-vl-1b-v2" @@ -483,3 +484,56 @@ async def test_nvidia_nim_ranking_endpoint_image_documents_and_top_n(): # top_n applied client-side on the converted response assert len(response.results) == 1 assert response.results[0]["index"] == 0 + + +class TestNvidiaNimRetrievalRerankRequestTransform: + """ + The default /v1/retrieval/{model}/reranking route keeps its existing + contract: top_n still maps to top_k, and structured documents now pass + through the same passage preservation as the /v1/ranking route. + """ + + def _build_request(self, documents, top_n=None): + config = NvidiaNimRerankConfig() + optional_params = config.map_cohere_rerank_params( + non_default_params=None, + model="nvidia/llama-3_2-nv-rerankqa-1b-v2", + drop_params=False, + query="which passage shows a cat?", + documents=documents, + top_n=top_n, + ) + return config.transform_rerank_request( + model="nvidia/llama-3_2-nv-rerankqa-1b-v2", + optional_rerank_params=optional_params, + headers={}, + ) + + def test_top_n_still_maps_to_top_k(self): + request_data = self._build_request(["a", "b"], top_n=1) + assert request_data["top_k"] == 1 + assert "top_n" not in request_data + + def test_string_documents_unchanged(self): + request_data = self._build_request(["passage one", "passage two"]) + assert request_data["passages"] == [ + {"text": "passage one"}, + {"text": "passage two"}, + ] + + def test_text_object_documents_unchanged(self): + request_data = self._build_request([TEXT_DOC]) + assert request_data["passages"] == [TEXT_DOC] + + def test_image_object_documents_are_preserved(self): + request_data = self._build_request([IMAGE_DOC, TEXT_DOC]) + assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + + def test_mixed_text_image_documents_are_preserved(self): + request_data = self._build_request([MIXED_DOC]) + assert request_data["passages"] == [MIXED_DOC] + + def test_unsupported_dict_documents_are_stringified(self): + doc = {"title": "no supported fields here"} + request_data = self._build_request([doc]) + assert request_data["passages"] == [{"text": json.dumps(doc)}] From 1d7db7564a4d408a1d7d464dfb470682376ba24b Mon Sep 17 00:00:00 2001 From: atomic <5234009+atomic@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:07:14 -0700 Subject: [PATCH 017/245] fix(nvidia_nim): satisfy current lint gates --- .../rerank/ranking_transformation.py | 34 +++++++++++-------- .../llms/nvidia_nim/rerank/transformation.py | 10 +++--- 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 58f26b2eb75..1b211faf8b3 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -78,26 +78,26 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): def map_cohere_rerank_params( self, - non_default_params: dict | None, + non_default_params: dict | None, # mutable-ok: matches BaseRerankConfig's request contract model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, Any]], # mutable-ok: matches BaseRerankConfig's document contract custom_llm_provider: str | None = None, top_n: int | None = None, - rank_fields: list[str] | None = None, + rank_fields: list[str] | None = None, # mutable-ok: matches BaseRerankConfig's field contract return_documents: bool | None = True, max_chunks_per_doc: int | None = None, max_tokens_per_doc: int | None = None, instruction: str | None = None, - ) -> dict: + ) -> dict: # mutable-ok: LiteLLM provider transforms return mutable request dictionaries """ Keep Cohere's top_n as-is instead of mapping it to top_k. The native /v1/ranking endpoint rejects top_k, so top_n is applied client-side after the response is converted. """ - optional_params = super().map_cohere_rerank_params( + optional_params: Final = super().map_cohere_rerank_params( non_default_params=non_default_params, model=model, drop_params=drop_params, @@ -132,14 +132,14 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): truncate. top_n is stashed and applied client-side in transform_rerank_response. """ - top_n = optional_rerank_params.get("top_n") + top_n: Final = optional_rerank_params.get("top_n") if top_n is not None: if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n < 1: raise ValueError(f"top_n must be a positive integer, got: {top_n!r}") self._client_side_top_n = top_n clean_model: Final = self._get_clean_model_name(model) - filtered_params: Final = { + filtered_params: Final = { # mutable-ok: the base transformer requires a mutable request dictionary k: v for k, v in optional_rerank_params.items() if k not in ("top_n", "top_k") } return super().transform_rerank_request( @@ -156,9 +156,9 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): model_response: RerankResponse, logging_obj: LiteLLMLoggingObj, api_key: str | None = None, - request_data: dict = {}, - optional_params: dict = {}, - litellm_params: dict = {}, + request_data: dict | None = None, # mutable-ok: matches BaseRerankConfig's response contract + optional_params: dict | None = None, # mutable-ok: matches BaseRerankConfig's response contract + litellm_params: dict | None = None, # mutable-ok: matches BaseRerankConfig's response contract ) -> RerankResponse: """ Convert the native ranking response, then apply top_n client-side. @@ -166,18 +166,22 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): /v1/ranking returns rankings sorted by relevance, but sort before truncating in case a server returns them unsorted. """ - response = super().transform_rerank_response( + resolved_request_data: Final = request_data or {} # mutable-ok: the base transformer requires a dictionary + resolved_optional_params: Final = optional_params or {} # mutable-ok: response options are keyed lookups + resolved_litellm_params: Final = litellm_params or {} # mutable-ok: the base transformer requires a dictionary + + response: Final = super().transform_rerank_response( model=model, raw_response=raw_response, model_response=model_response, logging_obj=logging_obj, api_key=api_key, - request_data=request_data, - optional_params=optional_params, - litellm_params=litellm_params, + request_data=resolved_request_data, + optional_params=resolved_optional_params, + litellm_params=resolved_litellm_params, ) - top_n = optional_params.get("top_n") or self._client_side_top_n + top_n: Final = resolved_optional_params.get("top_n") or self._client_side_top_n if top_n is not None and response.results is not None and len(response.results) > top_n: response.results = sorted( response.results, diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 7d4ecdb4cdc..65cb818295a 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -214,11 +214,11 @@ class NvidiaNimRerankConfig(BaseRerankConfig): elif isinstance(doc, dict): # Preserve structured passages (text, image, or mixed) so # VL rerank models receive image passages intact - supported_fields: NvidiaNimPassageObject = { - field: doc[field] # type: ignore[misc] - for field in self.SUPPORTED_PASSAGE_FIELDS - if field in doc - } + supported_fields: NvidiaNimPassageObject = {} # mutable-ok: assembling a request TypedDict + if "text" in doc: + supported_fields["text"] = doc["text"] + if "image" in doc: + supported_fields["image"] = doc["image"] if supported_fields: passages.append(supported_fields) else: From f3640447906f07c5bb0dc59ab454b0ac6d98f665 Mon Sep 17 00:00:00 2001 From: atomic <5234009+atomic@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:14:30 -0700 Subject: [PATCH 018/245] fix(nvidia_nim): scope image passages to ranking route --- .../nvidia_nim/rerank/ranking_transformation.py | 2 ++ litellm/llms/nvidia_nim/rerank/transformation.py | 15 +++++++-------- tests/llm_translation/test_nvidia_nim.py | 11 +++++++---- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 1b211faf8b3..f3b7ef7d8b0 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -37,6 +37,8 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): }' """ + SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text", "image") + def __init__(self) -> None: super().__init__() # top_n captured in transform_rerank_request and applied in diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 65cb818295a..bb07f9ec74f 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -54,10 +54,9 @@ class NvidiaNimRerankConfig(BaseRerankConfig): DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com" - # Structured document fields forwarded to the ranking API as-is. - # VL rerank models (e.g. nvidia/llama-nemotron-rerank-vl-1b-v2) accept - # image passages alongside text passages. - SUPPORTED_PASSAGE_FIELDS = ("text", "image") + # The legacy retrieval rerank route accepts text passages only. The native + # ranking subclass expands this tuple for VL models that accept images. + SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text",) def __init__(self) -> None: pass @@ -212,12 +211,12 @@ class NvidiaNimRerankConfig(BaseRerankConfig): if isinstance(doc, str): passages.append({"text": doc}) elif isinstance(doc, dict): - # Preserve structured passages (text, image, or mixed) so - # VL rerank models receive image passages intact + # Preserve only the structured passage fields supported by the + # selected rerank route. supported_fields: NvidiaNimPassageObject = {} # mutable-ok: assembling a request TypedDict - if "text" in doc: + if "text" in self.SUPPORTED_PASSAGE_FIELDS and "text" in doc: supported_fields["text"] = doc["text"] - if "image" in doc: + if "image" in self.SUPPORTED_PASSAGE_FIELDS and "image" in doc: supported_fields["image"] = doc["image"] if supported_fields: passages.append(supported_fields) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 5448fba4d0b..4bad94e0834 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -525,13 +525,16 @@ class TestNvidiaNimRetrievalRerankRequestTransform: request_data = self._build_request([TEXT_DOC]) assert request_data["passages"] == [TEXT_DOC] - def test_image_object_documents_are_preserved(self): + def test_image_object_documents_keep_retrieval_behavior(self): request_data = self._build_request([IMAGE_DOC, TEXT_DOC]) - assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + assert request_data["passages"] == [ + {"text": json.dumps(IMAGE_DOC)}, + TEXT_DOC, + ] - def test_mixed_text_image_documents_are_preserved(self): + def test_mixed_text_image_documents_keep_text_only(self): request_data = self._build_request([MIXED_DOC]) - assert request_data["passages"] == [MIXED_DOC] + assert request_data["passages"] == [{"text": MIXED_DOC["text"]}] def test_unsupported_dict_documents_are_stringified(self): doc = {"title": "no supported fields here"} From 25ad7dcb414a341f5d3e2013131d6b0d9a9e3acd Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Thu, 18 Jun 2026 23:59:59 -0700 Subject: [PATCH 019/245] fix(xai): bill web_search from server_side_tool_usage_details Use usage.server_side_tool_usage_details.web_search_calls at $5/1k calls instead of legacy num_sources_used/web_search_requests. Preserve tool usage details through Responses usage transform for accurate response cost. --- litellm/llms/xai/chat/transformation.py | 26 +++--- litellm/llms/xai/cost_calculator.py | 42 ++++------ litellm/responses/utils.py | 11 +++ .../llms/xai/test_xai_cost_calculator.py | 84 ++++--------------- 4 files changed, 56 insertions(+), 107 deletions(-) diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 9d06b609752..e9e9f205f94 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -18,7 +18,6 @@ from litellm.types.utils import ( Choices, ModelResponse, ModelResponseStream, - PromptTokensDetailsWrapper, Usage, ) @@ -248,7 +247,7 @@ class XAIChatConfig(OpenAIGPTConfig): XAI API returns empty string for finish_reason when using tools, so we need to fix this after the standard OpenAI transformation. - Also handles X.AI web search usage tracking by extracting num_sources_used. + Also handles X.AI web search usage tracking. """ # First, let the parent class handle the standard transformation @@ -351,25 +350,20 @@ class XAIChatConfig(OpenAIGPTConfig): def _enhance_usage_with_xai_web_search_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ - Extract num_sources_used from X.AI response and map it to web_search_requests. + Copy usage.server_side_tool_usage_details from the provider usage block + onto model_response.usage for tool cost calculation. """ if not hasattr(model_response, "usage") or model_response.usage is None: return usage: Final[Usage] = model_response.usage - num_sources_used = None - response_usage: Final = raw_response_json.get("usage", {}) - if isinstance(response_usage, dict) and "num_sources_used" in response_usage: - num_sources_used = response_usage.get("num_sources_used") - - # Map num_sources_used to web_search_requests for cost detection - if num_sources_used is not None and num_sources_used > 0: - if usage.prompt_tokens_details is None: - usage.prompt_tokens_details = PromptTokensDetailsWrapper() - - usage.prompt_tokens_details.web_search_requests = int(num_sources_used) - setattr(usage, "num_sources_used", int(num_sources_used)) - verbose_logger.debug("X.AI web search sources used: %s", num_sources_used) + response_usage: Final = raw_response_json.get("usage") + if not isinstance(response_usage, dict): + return + details = response_usage.get("server_side_tool_usage_details") + if details is not None: + setattr(usage, "server_side_tool_usage_details", details) + verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details) @staticmethod def _normalize_openai_compatible_usage_totals( diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 384388f3300..dcc96625975 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -4,6 +4,7 @@ Helper util for handling XAI-specific cost calculation - Handles XAI-specific reasoning token billing (billed as part of completion tokens) """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Final from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -12,6 +13,9 @@ from litellm.types.utils import Usage if TYPE_CHECKING: from litellm.types.utils import ModelInfo +# https://docs.x.ai/developers/pricing#tools-pricing +_WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ @@ -56,29 +60,17 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa """ Calculate the cost of web search requests for X.AI models. - X.AI Live Search costs $25 per 1,000 sources used. - Each source costs $0.025. - - The number of sources is stored in prompt_tokens_details.web_search_requests - by the transformation layer to be compatible with the existing detection system. + Uses usage.server_side_tool_usage_details.web_search_calls at $5 / 1k calls + (xAI tools pricing), not legacy num_sources_used / web_search_requests. """ - # Cost per source used: $25 per 1,000 sources = $0.025 per source - cost_per_source: Final = 25.0 / 1000.0 # $0.025 - - num_sources_used = 0 - - if ( - hasattr(usage, "prompt_tokens_details") - and usage.prompt_tokens_details is not None - and hasattr(usage.prompt_tokens_details, "web_search_requests") - and usage.prompt_tokens_details.web_search_requests is not None - ): - num_sources_used = int(usage.prompt_tokens_details.web_search_requests) - - # Fallback: try to get from num_sources_used if set directly - elif hasattr(usage, "num_sources_used") and usage.num_sources_used is not None: - num_sources_used = int(usage.num_sources_used) - - total_cost: Final = cost_per_source * num_sources_used - - return total_cost + _ = model_info + details = getattr(usage, "server_side_tool_usage_details", None) + if not isinstance(details, Mapping): + return 0.0 + try: + web_search_calls = int(details.get("web_search_calls") or 0) + except (TypeError, ValueError): + return 0.0 + if web_search_calls <= 0: + return 0.0 + return _WEB_SEARCH_COST_PER_CALL * web_search_calls diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index db2e515609c..659edb6db3b 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1089,12 +1089,23 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(output_tokens_details, "audio_tokens", None), ) + usage_kwargs: dict[str, Any] = {} + # Keep xAI tool billing fields; dropped if we only pass token fields below. + if isinstance(usage_input, dict): + if usage_input.get("server_side_tool_usage_details") is not None: + usage_kwargs["server_side_tool_usage_details"] = usage_input["server_side_tool_usage_details"] + else: + details = getattr(response_api_usage, "server_side_tool_usage_details", None) + if details is not None: + usage_kwargs["server_side_tool_usage_details"] = details + chat_usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, prompt_tokens_details=prompt_tokens_details, completion_tokens_details=completion_tokens_details, + **usage_kwargs, ) # Preserve cost attribute if it exists on ResponseAPIUsage diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 02fe7c8e68f..b166778deaf 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -9,7 +9,6 @@ import sys import litellm from litellm.types.utils import ( CompletionTokensDetailsWrapper, - PromptTokensDetailsWrapper, Usage, ) @@ -354,75 +353,28 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - def test_web_search_cost_calculation(self): - """Test web search cost calculation for X.AI models.""" - # Test with web_search_requests in prompt_tokens_details (primary path) - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=100, - web_search_requests=3, # 3 sources used - ), + def test_web_search_cost_via_server_side_tool_usage_details(self): + """usage.server_side_tool_usage_details.web_search_calls at $5/1k.""" + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + setattr( + usage, + "server_side_tool_usage_details", + { + "web_search_calls": 3, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + }, ) web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) + assert math.isclose(web_search_cost, 3 * (5.0 / 1000.0), rel_tol=1e-10) - # Expected cost: 3 sources * $0.025 per source = $0.075 - expected_cost = 3 * (25.0 / 1000.0) # 3 * $0.025 - - assert math.isclose(web_search_cost, expected_cost, rel_tol=1e-10) - assert math.isclose(web_search_cost, 0.075, rel_tol=1e-10) - - def test_web_search_cost_fallback_calculation(self): - """Test web search cost calculation using fallback num_sources_used.""" - # Test fallback: num_sources_used on usage object - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, - ) - # Manually set num_sources_used (as done by transformation layer) - setattr(usage, "num_sources_used", 5) - - web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) - - # Expected cost: 5 sources * $0.025 per source = $0.125 - expected_cost = 5 * (25.0 / 1000.0) # 5 * $0.025 - - assert math.isclose(web_search_cost, expected_cost, rel_tol=1e-10) - assert math.isclose(web_search_cost, 0.125, rel_tol=1e-10) - - def test_web_search_no_sources_used(self): - """Test web search cost calculation when no sources are used.""" - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=100, - web_search_requests=0, # No web search - ), - ) - - web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) - - # Expected cost: 0 sources * $0.025 per source = $0.0 - assert web_search_cost == 0.0 - - def test_web_search_cost_without_prompt_tokens_details(self): - """Test web search cost calculation when prompt_tokens_details is None.""" - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, - ) - - web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) - - # Expected cost: No web search data = $0.0 - assert web_search_cost == 0.0 + def test_web_search_cost_zero_without_details(self): + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" From 6963cfe0474d3676f192354d580608165aac91aa Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:26:37 -0700 Subject: [PATCH 020/245] style: apply black formatting to responses/utils.py --- litellm/responses/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 659edb6db3b..f9677e78ee7 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1095,7 +1095,9 @@ class ResponseAPILoggingUtils: if usage_input.get("server_side_tool_usage_details") is not None: usage_kwargs["server_side_tool_usage_details"] = usage_input["server_side_tool_usage_details"] else: - details = getattr(response_api_usage, "server_side_tool_usage_details", None) + details = getattr( + response_api_usage, "server_side_tool_usage_details", None + ) if details is not None: usage_kwargs["server_side_tool_usage_details"] = details From 014d59f4c4481747c8bdb825dffe4484b0cc8166 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:30:53 -0700 Subject: [PATCH 021/245] refactor(responses): pass through extra usage fields generically Avoid hard-coding provider-specific usage keys in shared Responses utilities; forward any non-standard usage attributes onto chat Usage for provider cost tracking (e.g. server_side_tool_usage_details). --- litellm/responses/utils.py | 56 +++++++++++++++---- .../responses/test_responses_utils.py | 17 ++++++ 2 files changed, 62 insertions(+), 11 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index f9677e78ee7..5e1412ee86b 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1022,6 +1022,20 @@ class ResponsesAPIRequestUtils: class ResponseAPILoggingUtils: + # Standard Responses usage keys mapped explicitly below; extras pass through to Usage. + _RESPONSE_API_USAGE_MAPPED_KEYS = frozenset( + { + "input_tokens", + "output_tokens", + "total_tokens", + "input_tokens_details", + "output_tokens_details", + "input_token_details", + "output_token_details", + "cost", # handled separately after Usage construction + } + ) + @staticmethod def _is_response_api_usage(usage: dict | ResponseAPIUsage) -> bool: """returns True if usage is from OpenAI Response API""" @@ -1031,6 +1045,33 @@ class ResponseAPILoggingUtils: return True return False + @staticmethod + def _extra_fields_from_response_api_usage( + usage_input: dict | ResponseAPIUsage, + response_api_usage: ResponseAPIUsage, + ) -> dict[str, Any]: + """ + Preserve provider/extension usage fields not part of the standard token mapping. + + ResponseAPIUsage allows extra attributes; without forwarding them, the rebuilt + chat Usage would drop fields needed for provider-specific cost tracking. + """ + extras: dict[str, Any] = {} + if isinstance(usage_input, dict): + for key, value in usage_input.items(): + if key not in ResponseAPILoggingUtils._RESPONSE_API_USAGE_MAPPED_KEYS and value is not None: + extras[key] = value + return extras + + model_extra = getattr(response_api_usage, "model_extra", None) or getattr( + response_api_usage, "__pydantic_extra__", None + ) + if isinstance(model_extra, dict): + for key, value in model_extra.items(): + if key not in ResponseAPILoggingUtils._RESPONSE_API_USAGE_MAPPED_KEYS and value is not None: + extras[key] = value + return extras + @staticmethod def _transform_response_api_usage_to_chat_usage( usage_input: dict | ResponseAPIUsage | None, @@ -1089,17 +1130,10 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(output_tokens_details, "audio_tokens", None), ) - usage_kwargs: dict[str, Any] = {} - # Keep xAI tool billing fields; dropped if we only pass token fields below. - if isinstance(usage_input, dict): - if usage_input.get("server_side_tool_usage_details") is not None: - usage_kwargs["server_side_tool_usage_details"] = usage_input["server_side_tool_usage_details"] - else: - details = getattr( - response_api_usage, "server_side_tool_usage_details", None - ) - if details is not None: - usage_kwargs["server_side_tool_usage_details"] = details + usage_kwargs: Final = ResponseAPILoggingUtils._extra_fields_from_response_api_usage( + usage_input=usage_input, + response_api_usage=response_api_usage, + ) chat_usage: Final = Usage( prompt_tokens=prompt_tokens, diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 0141cf5d96a..8785ff1da43 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -278,6 +278,23 @@ class TestResponseAPILoggingUtils: and result.prompt_tokens_details.cached_tokens == 2 ) + def test_transform_response_api_usage_preserves_extra_usage_fields(self): + """Non-standard usage keys pass through for provider cost tracking.""" + usage = { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "server_side_tool_usage_details": {"web_search_calls": 2}, + "num_server_side_tools_used": 2, + } + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + assert getattr(result, "server_side_tool_usage_details", None) == { + "web_search_calls": 2 + } + assert getattr(result, "num_server_side_tools_used", None) == 2 + def test_transform_response_api_usage_with_none_values(self): """Test transformation handles None values properly""" # Setup From 3aea951e6cce99787cd8e7d1bdf73b4a5aa4d46c Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:36:43 -0700 Subject: [PATCH 022/245] refactor(xai): keep Responses tool usage pass-through in llms/xai Revert shared responses/utils.py extras forwarding. Attach server_side_tool_usage_details on chat Usage inside XAIResponsesAPIConfig so cost calc keeps web_search_calls without provider logic in shared utils. --- litellm/llms/xai/responses/transformation.py | 50 ++++++++++++++++++- litellm/responses/utils.py | 47 ----------------- .../responses/test_responses_utils.py | 17 ------- 3 files changed, 48 insertions(+), 66 deletions(-) diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 48fb95d9411..fbff9876d49 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,16 +1,22 @@ from typing import TYPE_CHECKING, Any, Final +import httpx + import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.responses.utils import ResponseAPILoggingUtils from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ( + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, +) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LlmProviders +from litellm.types.utils import LlmProviders, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -51,6 +57,46 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return supported_params + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + """ + Attach xAI tool usage details onto a chat Usage object. + + Cost calculation normalizes Responses usage via + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage, which + drops non-standard fields unless usage is already a chat Usage instance. + """ + response = super().transform_response_api_response( + model=model, raw_response=raw_response, logging_obj=logging_obj + ) + self._attach_server_side_tool_usage_details_to_usage(response) + return response + + @staticmethod + def _attach_server_side_tool_usage_details_to_usage( + response: ResponsesAPIResponse, + ) -> None: + if response.usage is None: + return + + details = getattr(response.usage, "server_side_tool_usage_details", None) + if details is None and isinstance(response.usage, dict): + details = response.usage.get("server_side_tool_usage_details") + if details is None: + return + + if isinstance(response.usage, Usage): + setattr(response.usage, "server_side_tool_usage_details", details) + return + + chat_usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response.usage) + setattr(chat_usage, "server_side_tool_usage_details", details) + response.usage = chat_usage # type: ignore[assignment] + def _transform_web_search_tool(self, tool: dict[str, Any]) -> XAIWebSearchTool | dict[str, Any]: """ Transform web_search tool to XAI format. diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 5e1412ee86b..db2e515609c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1022,20 +1022,6 @@ class ResponsesAPIRequestUtils: class ResponseAPILoggingUtils: - # Standard Responses usage keys mapped explicitly below; extras pass through to Usage. - _RESPONSE_API_USAGE_MAPPED_KEYS = frozenset( - { - "input_tokens", - "output_tokens", - "total_tokens", - "input_tokens_details", - "output_tokens_details", - "input_token_details", - "output_token_details", - "cost", # handled separately after Usage construction - } - ) - @staticmethod def _is_response_api_usage(usage: dict | ResponseAPIUsage) -> bool: """returns True if usage is from OpenAI Response API""" @@ -1045,33 +1031,6 @@ class ResponseAPILoggingUtils: return True return False - @staticmethod - def _extra_fields_from_response_api_usage( - usage_input: dict | ResponseAPIUsage, - response_api_usage: ResponseAPIUsage, - ) -> dict[str, Any]: - """ - Preserve provider/extension usage fields not part of the standard token mapping. - - ResponseAPIUsage allows extra attributes; without forwarding them, the rebuilt - chat Usage would drop fields needed for provider-specific cost tracking. - """ - extras: dict[str, Any] = {} - if isinstance(usage_input, dict): - for key, value in usage_input.items(): - if key not in ResponseAPILoggingUtils._RESPONSE_API_USAGE_MAPPED_KEYS and value is not None: - extras[key] = value - return extras - - model_extra = getattr(response_api_usage, "model_extra", None) or getattr( - response_api_usage, "__pydantic_extra__", None - ) - if isinstance(model_extra, dict): - for key, value in model_extra.items(): - if key not in ResponseAPILoggingUtils._RESPONSE_API_USAGE_MAPPED_KEYS and value is not None: - extras[key] = value - return extras - @staticmethod def _transform_response_api_usage_to_chat_usage( usage_input: dict | ResponseAPIUsage | None, @@ -1130,18 +1089,12 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(output_tokens_details, "audio_tokens", None), ) - usage_kwargs: Final = ResponseAPILoggingUtils._extra_fields_from_response_api_usage( - usage_input=usage_input, - response_api_usage=response_api_usage, - ) - chat_usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, prompt_tokens_details=prompt_tokens_details, completion_tokens_details=completion_tokens_details, - **usage_kwargs, ) # Preserve cost attribute if it exists on ResponseAPIUsage diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 8785ff1da43..0141cf5d96a 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -278,23 +278,6 @@ class TestResponseAPILoggingUtils: and result.prompt_tokens_details.cached_tokens == 2 ) - def test_transform_response_api_usage_preserves_extra_usage_fields(self): - """Non-standard usage keys pass through for provider cost tracking.""" - usage = { - "input_tokens": 10, - "output_tokens": 5, - "total_tokens": 15, - "server_side_tool_usage_details": {"web_search_calls": 2}, - "num_server_side_tools_used": 2, - } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) - assert getattr(result, "server_side_tool_usage_details", None) == { - "web_search_calls": 2 - } - assert getattr(result, "num_server_side_tools_used", None) == 2 - def test_transform_response_api_usage_with_none_values(self): """Test transformation handles None values properly""" # Setup From ea98d8e116af74809aab1149ad34003fba0f516f Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:43:15 -0700 Subject: [PATCH 023/245] fix(xai): read tool usage details from ResponseAPIUsage extras Also inspect model_extra when attaching server_side_tool_usage_details for Responses cost tracking. --- litellm/llms/xai/responses/transformation.py | 22 +++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index fbff9876d49..7cf2f48f956 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -76,6 +76,22 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): self._attach_server_side_tool_usage_details_to_usage(response) return response + @staticmethod + def _server_side_tool_usage_details_from_usage(usage: Any) -> Any: + if usage is None: + return None + if isinstance(usage, dict): + return usage.get("server_side_tool_usage_details") + details = getattr(usage, "server_side_tool_usage_details", None) + if details is not None: + return details + model_extra = getattr(usage, "model_extra", None) or getattr( + usage, "__pydantic_extra__", None + ) + if isinstance(model_extra, dict): + return model_extra.get("server_side_tool_usage_details") + return None + @staticmethod def _attach_server_side_tool_usage_details_to_usage( response: ResponsesAPIResponse, @@ -83,9 +99,9 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): if response.usage is None: return - details = getattr(response.usage, "server_side_tool_usage_details", None) - if details is None and isinstance(response.usage, dict): - details = response.usage.get("server_side_tool_usage_details") + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( + response.usage + ) if details is None: return From 8687d7372ad3f2f40586fa40102b7571b678de52 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:55:35 -0700 Subject: [PATCH 024/245] fix(xai): gate web search cost on server_side_tool_usage_details Treat positive web_search_calls as a web-search signal in built-in tool cost gating, and mirror counts onto prompt_tokens_details.web_search_requests when attaching xAI tool usage details so charges are not skipped. --- .../llm_cost_calc/tool_call_cost_tracking.py | 18 +++++++ litellm/llms/xai/chat/transformation.py | 5 +- litellm/llms/xai/cost_calculator.py | 25 +++++++++- litellm/llms/xai/responses/transformation.py | 7 ++- .../llms/xai/test_xai_cost_calculator.py | 50 ++++++++++++++++++- 5 files changed, 99 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 3744be5bc79..eb9473438f1 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -311,6 +311,22 @@ class StandardBuiltInToolCostTracking: return Usage(server_tool_use=server_tool_use) return usage.model_copy(update={"server_tool_use": server_tool_use}) + @staticmethod + def _usage_has_server_side_web_search_calls(usage: Usage | None) -> bool: + """True when usage.server_side_tool_usage_details.web_search_calls > 0.""" + if usage is None: + return False + details = getattr(usage, "server_side_tool_usage_details", None) + if details is None: + return False + try: + web_search_calls = ( + details.get("web_search_calls") if isinstance(details, dict) else getattr(details, "web_search_calls", None) + ) + return int(web_search_calls or 0) > 0 + except (TypeError, ValueError): + return False + @staticmethod def response_object_includes_web_search_call(response_object: Any, usage: Usage | None = None) -> bool: """ @@ -328,6 +344,8 @@ class StandardBuiltInToolCostTracking: if get_anthropic_web_search_requests_from_response(response_object) is not None: return True + if StandardBuiltInToolCostTracking._usage_has_server_side_web_search_calls(usage): + return True if isinstance(response_object, ModelResponse): # chat completions only include url_citation annotations when a web search call is made diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index e9e9f205f94..1129872163c 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -12,6 +12,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_messages, ) from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.cost_calculator import ( + apply_server_side_tool_usage_details_to_usage, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( @@ -362,7 +365,7 @@ class XAIChatConfig(OpenAIGPTConfig): return details = response_usage.get("server_side_tool_usage_details") if details is not None: - setattr(usage, "server_side_tool_usage_details", details) + apply_server_side_tool_usage_details_to_usage(usage, details) verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details) @staticmethod diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index dcc96625975..981cdb7ac88 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -5,10 +5,10 @@ Helper util for handling XAI-specific cost calculation """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Any, Final from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import Usage +from litellm.types.utils import PromptTokensDetailsWrapper, Usage if TYPE_CHECKING: from litellm.types.utils import ModelInfo @@ -17,6 +17,27 @@ if TYPE_CHECKING: _WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 +def apply_server_side_tool_usage_details_to_usage( + usage: Usage, details: Mapping[str, Any] | None +) -> None: + """ + Attach server_side_tool_usage_details and mirror web_search_calls onto + prompt_tokens_details.web_search_requests for built-in tool cost gating. + """ + if details is None: + return + setattr(usage, "server_side_tool_usage_details", details) + try: + web_search_calls = int(details.get("web_search_calls") or 0) + except (TypeError, ValueError): + return + if web_search_calls <= 0: + return + if usage.prompt_tokens_details is None: + usage.prompt_tokens_details = PromptTokensDetailsWrapper() + usage.prompt_tokens_details.web_search_requests = web_search_calls + + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 7cf2f48f956..983f20fcf03 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -8,6 +8,9 @@ from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.cost_calculator import ( + apply_server_side_tool_usage_details_to_usage, +) from litellm.responses.utils import ResponseAPILoggingUtils from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -106,11 +109,11 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return if isinstance(response.usage, Usage): - setattr(response.usage, "server_side_tool_usage_details", details) + apply_server_side_tool_usage_details_to_usage(response.usage, details) return chat_usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response.usage) - setattr(chat_usage, "server_side_tool_usage_details", details) + apply_server_side_tool_usage_details_to_usage(chat_usage, details) response.usage = chat_usage # type: ignore[assignment] def _transform_web_search_tool(self, tool: dict[str, Any]) -> XAIWebSearchTool | dict[str, Any]: diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index b166778deaf..6435cf6218b 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -16,7 +16,15 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.llms.xai.cost_calculator import cost_per_token, cost_per_web_search_request +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, +) +from litellm.llms.xai.cost_calculator import ( + apply_server_side_tool_usage_details_to_usage, + cost_per_token, + cost_per_web_search_request, +) +from litellm.types.llms.openai import ResponsesAPIResponse class TestXAICostCalculator: @@ -376,6 +384,46 @@ class TestXAICostCalculator: usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + def test_apply_details_sets_web_search_requests_for_cost_gate(self): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + apply_server_side_tool_usage_details_to_usage( + usage, {"web_search_calls": 2, "x_search_calls": 0} + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.web_search_requests == 2 + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=object(), usage=usage + ) + + def test_gate_detects_server_side_tool_usage_details_without_web_search_output( + self, + ): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + setattr( + usage, + "server_side_tool_usage_details", + {"web_search_calls": 1}, + ) + response = ResponsesAPIResponse.model_construct( + id="resp_test", + created_at=0, + output=[{"type": "message", "role": "assistant", "content": []}], + usage=None, + ) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage + ) + assert ( + StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="grok-4.3", + response_object=response, + usage=usage, + standard_built_in_tools_params={}, + custom_llm_provider="xai", + ) + == 5.0 / 1000.0 + ) + def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) From 7aaa9358aa67e5060abce392140947218bc928b9 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 23:05:35 -0700 Subject: [PATCH 025/245] fix(xai): read web_search per-call rate from model_info Use search_context_cost_per_query from the model cost map (with $5/1k fallback) so web search billing can change via pricing JSON updates. --- litellm/llms/xai/cost_calculator.py | 38 ++++++++++++++++--- .../llms/xai/test_xai_cost_calculator.py | 15 +++++++- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 981cdb7ac88..ebeac61fb43 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -13,8 +13,8 @@ from litellm.types.utils import PromptTokensDetailsWrapper, Usage if TYPE_CHECKING: from litellm.types.utils import ModelInfo -# https://docs.x.ai/developers/pricing#tools-pricing -_WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 +# https://docs.x.ai/developers/pricing#tools-pricing — default when unset in model map +_DEFAULT_WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 def apply_server_side_tool_usage_details_to_usage( @@ -77,14 +77,40 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: return prompt_cost, completion_cost +def _web_search_cost_per_call_from_model_info(model_info: "ModelInfo") -> float: + """ + Per-invocation web_search price from model_info when configured. + + Prefer ``search_context_cost_per_query`` (same shape as Gemini/Anthropic web + search pricing in the model cost map). Fall back to current xAI list pricing. + """ + search_costs = model_info.get("search_context_cost_per_query") or {} + if isinstance(search_costs, Mapping): + for key in ( + "search_context_size_medium", + "search_context_size_low", + "search_context_size_high", + ): + value = search_costs.get(key) + if value is None: + continue + try: + cost = float(value) + except (TypeError, ValueError): + continue + if cost > 0: + return cost + return _DEFAULT_WEB_SEARCH_COST_PER_CALL + + def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float: """ Calculate the cost of web search requests for X.AI models. - Uses usage.server_side_tool_usage_details.web_search_calls at $5 / 1k calls - (xAI tools pricing), not legacy num_sources_used / web_search_requests. + Counts invocations from usage.server_side_tool_usage_details.web_search_calls. + Per-call rate comes from model_info.search_context_cost_per_query when set, + otherwise the default xAI tools rate ($5 / 1k calls). """ - _ = model_info details = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): return 0.0 @@ -94,4 +120,4 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa return 0.0 if web_search_calls <= 0: return 0.0 - return _WEB_SEARCH_COST_PER_CALL * web_search_calls + return _web_search_cost_per_call_from_model_info(model_info) * web_search_calls diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 6435cf6218b..585a23725e6 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -362,7 +362,7 @@ class TestXAICostCalculator: assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) def test_web_search_cost_via_server_side_tool_usage_details(self): - """usage.server_side_tool_usage_details.web_search_calls at $5/1k.""" + """usage.server_side_tool_usage_details.web_search_calls at default $5/1k.""" usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) setattr( usage, @@ -380,6 +380,19 @@ class TestXAICostCalculator: web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) assert math.isclose(web_search_cost, 3 * (5.0 / 1000.0), rel_tol=1e-10) + def test_web_search_cost_uses_model_info_search_context_pricing(self): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 2}) + model_info = { + "search_context_cost_per_query": { + "search_context_size_medium": 0.01, + } + } + web_search_cost = cost_per_web_search_request( + usage=usage, model_info=model_info + ) + assert math.isclose(web_search_cost, 0.02, rel_tol=1e-10) + def test_web_search_cost_zero_without_details(self): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 From ea493543d7094b6af5b208ea26e3b55731adaf42 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 23:32:13 -0700 Subject: [PATCH 026/245] fix(xai): attach tool usage details on Responses stream terminal events Apply server_side_tool_usage_details on completed/incomplete/failed streaming events so stream=true web_search is billed like non-stream. --- litellm/llms/xai/responses/transformation.py | 29 ++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 983f20fcf03..de041854b3a 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -14,8 +14,12 @@ from litellm.llms.xai.cost_calculator import ( from litellm.responses.utils import ResponseAPILoggingUtils from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, + ResponsesAPIStreamingResponse, ) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams @@ -79,6 +83,31 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): self._attach_server_side_tool_usage_details_to_usage(response) return response + def transform_streaming_response( + self, + model: str, + parsed_chunk: dict, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIStreamingResponse: + """ + Preserve xAI tool usage on streaming terminal events for cost logging. + + Completed/incomplete/failed events embed a full ResponsesAPIResponse; without + attaching server_side_tool_usage_details here, stream=true web_search usage is + dropped when usage is normalized for billing. + """ + event = super().transform_streaming_response( + model=model, parsed_chunk=parsed_chunk, logging_obj=logging_obj + ) + if isinstance( + event, + (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent), + ): + embedded_response = getattr(event, "response", None) + if isinstance(embedded_response, ResponsesAPIResponse): + self._attach_server_side_tool_usage_details_to_usage(embedded_response) + return event + @staticmethod def _server_side_tool_usage_details_from_usage(usage: Any) -> Any: if usage is None: From c03a6076cebe07f574bdf11707728c2c61d019d7 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 20 Jun 2026 10:43:55 -0700 Subject: [PATCH 027/245] test(xai): cover Responses tool usage attach helpers Add unit tests for server_side_tool_usage_details extraction/attach and streaming completed-event pass-through in XAIResponsesAPIConfig. --- .../test_xai_responses_transformation.py | 157 +++++++++++++++++- 1 file changed, 155 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index fb98dc0a917..c309bd4a983 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -9,14 +9,22 @@ Source: litellm/llms/xai/responses/transformation.py import os import sys +from unittest.mock import MagicMock sys.path.insert(0, os.path.abspath("../../../../..")) import pytest from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams -from litellm.types.utils import LlmProviders +from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, +) +from litellm.types.utils import LlmProviders, Usage from litellm.utils import ProviderConfigManager @@ -309,3 +317,148 @@ class TestXAIResponsesAPITransformation: # Verify function tool is unchanged assert result["tools"][3]["type"] == "function" assert result["tools"][3]["name"] == "get_weather" + + +class TestXAIResponsesToolUsageAttach: + """Tests for server_side_tool_usage_details attach helpers (cost billing).""" + + _TOOL_DETAILS = { + "web_search_calls": 2, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + } + + def test_server_side_tool_usage_details_from_usage_dict(self): + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( + {"server_side_tool_usage_details": self._TOOL_DETAILS} + ) + assert details == self._TOOL_DETAILS + + def test_server_side_tool_usage_details_from_usage_attr(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr(usage, "server_side_tool_usage_details", self._TOOL_DETAILS) + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( + usage + ) + assert details == self._TOOL_DETAILS + + def test_server_side_tool_usage_details_from_model_extra(self): + usage = ResponseAPIUsage( + input_tokens=10, + output_tokens=5, + total_tokens=15, + server_side_tool_usage_details=self._TOOL_DETAILS, + ) + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( + usage + ) + assert details == self._TOOL_DETAILS + + def test_server_side_tool_usage_details_from_usage_none(self): + assert ( + XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage(None) + is None + ) + assert ( + XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( + Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1) + ) + is None + ) + + def test_attach_noop_when_usage_missing(self): + response = ResponsesAPIResponse.model_construct( + id="resp_1", created_at=0, output=[], usage=None + ) + XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) + assert response.usage is None + + def test_attach_noop_when_details_missing(self): + usage = ResponseAPIUsage(input_tokens=3, output_tokens=1, total_tokens=4) + response = ResponsesAPIResponse.model_construct( + id="resp_2", created_at=0, output=[], usage=usage + ) + XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) + assert isinstance(response.usage, ResponseAPIUsage) + + def test_attach_converts_response_api_usage_to_chat_usage(self): + usage = ResponseAPIUsage( + input_tokens=100, + output_tokens=20, + total_tokens=120, + server_side_tool_usage_details=self._TOOL_DETAILS, + ) + response = ResponsesAPIResponse.model_construct( + id="resp_3", created_at=0, output=[], usage=usage + ) + XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) + + assert isinstance(response.usage, Usage) + assert response.usage.prompt_tokens == 100 + assert response.usage.completion_tokens == 20 + assert getattr(response.usage, "server_side_tool_usage_details") == ( + self._TOOL_DETAILS + ) + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.web_search_requests == 2 + + def test_attach_updates_existing_chat_usage_in_place(self): + usage = Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10) + setattr(usage, "server_side_tool_usage_details", self._TOOL_DETAILS) + response = ResponsesAPIResponse.model_construct( + id="resp_4", created_at=0, output=[], usage=usage + ) + XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) + + assert response.usage is usage + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.web_search_requests == 2 + + def test_transform_streaming_response_completed_attaches_tool_usage(self): + config = XAIResponsesAPIConfig() + chunk = { + "type": "response.completed", + "response": { + "id": "resp_stream", + "created_at": 1, + "output": [], + "usage": { + "input_tokens": 50, + "output_tokens": 10, + "total_tokens": 60, + "server_side_tool_usage_details": self._TOOL_DETAILS, + }, + }, + } + event = config.transform_streaming_response( + model="grok-4.3", parsed_chunk=chunk, logging_obj=MagicMock() + ) + + assert isinstance(event, ResponseCompletedEvent) + assert isinstance(event.response.usage, Usage) + assert getattr(event.response.usage, "server_side_tool_usage_details") == ( + self._TOOL_DETAILS + ) + assert event.response.usage.prompt_tokens_details is not None + assert event.response.usage.prompt_tokens_details.web_search_requests == 2 + + def test_transform_streaming_response_non_terminal_event_unchanged(self): + config = XAIResponsesAPIConfig() + chunk = { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + } + event = config.transform_streaming_response( + model="grok-4.3", parsed_chunk=chunk, logging_obj=MagicMock() + ) + assert getattr(event, "type", None) is not None + assert not isinstance( + event, + (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent), + ) From 478118ac36bdaf48e3847ec281a12f79e09a1f2f Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 20 Jun 2026 10:46:37 -0700 Subject: [PATCH 028/245] test(xai): expand cost_calculator coverage for web search helpers Add unit tests for apply_server_side_tool_usage_details_to_usage edge cases and model_info-driven web_search per-call pricing fallbacks. --- .../llms/xai/test_xai_cost_calculator.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 585a23725e6..4e1c8cacfd7 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -9,6 +9,7 @@ import sys import litellm from litellm.types.utils import ( CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, Usage, ) @@ -20,6 +21,8 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.llms.xai.cost_calculator import ( + _DEFAULT_WEB_SEARCH_COST_PER_CALL, + _web_search_cost_per_call_from_model_info, apply_server_side_tool_usage_details_to_usage, cost_per_token, cost_per_web_search_request, @@ -484,3 +487,112 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + +class TestXAIWebSearchCostHelpers: + """Focused coverage for web_search / tool-usage helpers in cost_calculator.py.""" + + def test_apply_details_noop_when_details_none(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + apply_server_side_tool_usage_details_to_usage(usage, None) + assert getattr(usage, "server_side_tool_usage_details", None) is None + + def test_apply_details_sets_attr_but_skips_mirror_when_web_search_zero(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + details = {"web_search_calls": 0, "x_search_calls": 3} + apply_server_side_tool_usage_details_to_usage(usage, details) + assert getattr(usage, "server_side_tool_usage_details") == details + assert ( + usage.prompt_tokens_details is None + or usage.prompt_tokens_details.web_search_requests is None + ) + + def test_apply_details_skips_mirror_when_web_search_calls_invalid(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + details = {"web_search_calls": "not-a-number"} + apply_server_side_tool_usage_details_to_usage(usage, details) + assert getattr(usage, "server_side_tool_usage_details") == details + assert usage.prompt_tokens_details is None + + def test_apply_details_updates_existing_prompt_tokens_details(self): + usage = Usage( + prompt_tokens=1, + completion_tokens=1, + total_tokens=2, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=7), + ) + apply_server_side_tool_usage_details_to_usage(usage, {"web_search_calls": 4}) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 7 + assert usage.prompt_tokens_details.web_search_requests == 4 + + def test_web_search_cost_per_call_default_when_model_info_empty(self): + assert ( + _web_search_cost_per_call_from_model_info({}) + == _DEFAULT_WEB_SEARCH_COST_PER_CALL + ) + + def test_web_search_cost_per_call_prefers_medium_over_low(self): + model_info = { + "search_context_cost_per_query": { + "search_context_size_low": 0.001, + "search_context_size_medium": 0.009, + } + } + assert _web_search_cost_per_call_from_model_info(model_info) == 0.009 + + def test_web_search_cost_per_call_falls_back_to_low_then_high(self): + assert ( + _web_search_cost_per_call_from_model_info( + {"search_context_cost_per_query": {"search_context_size_low": 0.003}} + ) + == 0.003 + ) + assert ( + _web_search_cost_per_call_from_model_info( + {"search_context_cost_per_query": {"search_context_size_high": 0.007}} + ) + == 0.007 + ) + + def test_web_search_cost_per_call_ignores_zero_and_invalid_values(self): + assert ( + _web_search_cost_per_call_from_model_info( + { + "search_context_cost_per_query": { + "search_context_size_medium": 0, + "search_context_size_low": "bad", + } + } + ) + == _DEFAULT_WEB_SEARCH_COST_PER_CALL + ) + + def test_cost_per_web_search_request_zero_when_details_not_mapping(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr(usage, "server_side_tool_usage_details", "invalid") + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_cost_per_web_search_request_zero_when_web_search_calls_invalid(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr( + usage, + "server_side_tool_usage_details", + {"web_search_calls": object()}, + ) + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_cost_per_web_search_request_zero_when_web_search_calls_zero(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr( + usage, + "server_side_tool_usage_details", + {"web_search_calls": 0, "x_search_calls": 5}, + ) + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_cost_per_web_search_request_uses_default_rate_without_model_pricing(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 4}) + cost = cost_per_web_search_request(usage=usage, model_info={}) + assert math.isclose(cost, 4 * _DEFAULT_WEB_SEARCH_COST_PER_CALL, rel_tol=1e-10) From 8ad0a57387631706b293b7fc6168ed4374825d7d Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 20 Jun 2026 10:48:25 -0700 Subject: [PATCH 029/245] style: drop unused pytest import in xAI responses tests --- .../llms/xai/responses/test_xai_responses_transformation.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index c309bd4a983..b2072867539 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -13,8 +13,6 @@ from unittest.mock import MagicMock sys.path.insert(0, os.path.abspath("../../../../..")) -import pytest - from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig from litellm.types.llms.openai import ( ResponseAPIUsage, From 74100989a29144ae2ce32b0d173e68b3e00fb4ec Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 20 Jun 2026 10:57:36 -0700 Subject: [PATCH 030/245] revert: remove xAI-specific web search gate from shared cost tracking Gate web search like OpenAI (output/annotations/web_search_requests). xAI uses server_side_tool_usage_details only for per-call cost math, with web_search_requests mirrored in llms/xai for existing gate compatibility. --- .../llm_cost_calc/tool_call_cost_tracking.py | 18 ----------- .../llms/xai/test_xai_cost_calculator.py | 32 ------------------- 2 files changed, 50 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index eb9473438f1..3744be5bc79 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -311,22 +311,6 @@ class StandardBuiltInToolCostTracking: return Usage(server_tool_use=server_tool_use) return usage.model_copy(update={"server_tool_use": server_tool_use}) - @staticmethod - def _usage_has_server_side_web_search_calls(usage: Usage | None) -> bool: - """True when usage.server_side_tool_usage_details.web_search_calls > 0.""" - if usage is None: - return False - details = getattr(usage, "server_side_tool_usage_details", None) - if details is None: - return False - try: - web_search_calls = ( - details.get("web_search_calls") if isinstance(details, dict) else getattr(details, "web_search_calls", None) - ) - return int(web_search_calls or 0) > 0 - except (TypeError, ValueError): - return False - @staticmethod def response_object_includes_web_search_call(response_object: Any, usage: Usage | None = None) -> bool: """ @@ -344,8 +328,6 @@ class StandardBuiltInToolCostTracking: if get_anthropic_web_search_requests_from_response(response_object) is not None: return True - if StandardBuiltInToolCostTracking._usage_has_server_side_web_search_calls(usage): - return True if isinstance(response_object, ModelResponse): # chat completions only include url_citation annotations when a web search call is made diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 4e1c8cacfd7..1102a63c917 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -27,9 +27,6 @@ from litellm.llms.xai.cost_calculator import ( cost_per_token, cost_per_web_search_request, ) -from litellm.types.llms.openai import ResponsesAPIResponse - - class TestXAICostCalculator: """Test suite for XAI cost calculation functionality.""" @@ -411,35 +408,6 @@ class TestXAICostCalculator: response_object=object(), usage=usage ) - def test_gate_detects_server_side_tool_usage_details_without_web_search_output( - self, - ): - usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) - setattr( - usage, - "server_side_tool_usage_details", - {"web_search_calls": 1}, - ) - response = ResponsesAPIResponse.model_construct( - id="resp_test", - created_at=0, - output=[{"type": "message", "role": "assistant", "content": []}], - usage=None, - ) - assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( - response_object=response, usage=usage - ) - assert ( - StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( - model="grok-4.3", - response_object=response, - usage=usage, - standard_built_in_tools_params={}, - custom_llm_provider="xai", - ) - == 5.0 / 1000.0 - ) - def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) From a9277b4b6e248b8323c280f657107ed28660c2ba Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 20 Jun 2026 11:03:30 -0700 Subject: [PATCH 031/245] style: black format xAI cost calculator tests --- tests/test_litellm/llms/xai/test_xai_cost_calculator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 1102a63c917..80503514190 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -27,6 +27,8 @@ from litellm.llms.xai.cost_calculator import ( cost_per_token, cost_per_web_search_request, ) + + class TestXAICostCalculator: """Test suite for XAI cost calculation functionality.""" From 749a8b0701c630b054a237b19090292dd9bb2318 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 8 Aug 2026 16:49:09 -0700 Subject: [PATCH 032/245] fix(xai): keep chat Usage through Responses completions bridge xAI already converts Responses usage to chat Usage so web_search_calls survive cost tracking. The chat completions bridge then re-ran the Responses usage transform and crashed on missing input_tokens. Pass through already-chat Usage and chat-shaped dumps instead --- litellm/responses/utils.py | 10 +- .../test_xai_responses_transformation.py | 124 ++++++-------- .../responses/test_responses_utils.py | 151 +++++++++--------- 3 files changed, 137 insertions(+), 148 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index db2e515609c..27dd4230923 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1033,13 +1033,17 @@ class ResponseAPILoggingUtils: @staticmethod def _transform_response_api_usage_to_chat_usage( - usage_input: dict | ResponseAPIUsage | None, + usage_input: dict | ResponseAPIUsage | Usage | None, ) -> Usage: """ Transforms ResponseAPIUsage or ImageUsage to a Usage object. Both have the same spec with input_tokens, output_tokens, and input_tokens_details (text_tokens, image_tokens). + + Providers that already converted usage to chat Usage (e.g. xAI Responses + attaching server_side_tool_usage_details) are returned as-is so the chat + completions bridge can re-run this helper without dropping extra fields. """ if usage_input is None: return Usage( @@ -1047,6 +1051,10 @@ class ResponseAPILoggingUtils: completion_tokens=0, total_tokens=0, ) + if isinstance(usage_input, Usage): + return usage_input + if isinstance(usage_input, dict) and not ResponseAPILoggingUtils._is_response_api_usage(usage_input): + return Usage(**usage_input) response_api_usage: ResponseAPIUsage if isinstance(usage_input, dict): usage_input = dict(usage_input) # shallow copy; avoid mutating caller diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index b2072867539..e0e526fd38e 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -37,43 +37,29 @@ class TestXAIResponsesAPITransformation: ) assert config is not None, "Config should not be None for XAI provider" - assert isinstance( - config, XAIResponsesAPIConfig - ), f"Expected XAIResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.XAI - ), "custom_llm_provider should be XAI" + assert isinstance(config, XAIResponsesAPIConfig), f"Expected XAIResponsesAPIConfig, got {type(config)}" + assert config.custom_llm_provider == LlmProviders.XAI, "custom_llm_provider should be XAI" def test_code_interpreter_container_field_removed(self): """Test that container field is removed from code_interpreter tools""" config = XAIResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - tools=[{"type": "code_interpreter", "container": {"type": "auto"}}] - ) + params = ResponsesAPIOptionalRequestParams(tools=[{"type": "code_interpreter", "container": {"type": "auto"}}]) - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) + result = config.map_openai_params(response_api_optional_params=params, model="grok-4-fast", drop_params=False) assert "tools" in result assert len(result["tools"]) == 1 assert result["tools"][0]["type"] == "code_interpreter" - assert ( - "container" not in result["tools"][0] - ), "Container field should be removed" + assert "container" not in result["tools"][0], "Container field should be removed" def test_instructions_parameter_dropped(self): """Test that instructions parameter is dropped for XAI""" config = XAIResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - instructions="You are a helpful assistant.", temperature=0.7 - ) + params = ResponsesAPIOptionalRequestParams(instructions="You are a helpful assistant.", temperature=0.7) - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) + result = config.map_openai_params(response_api_optional_params=params, model="grok-4-fast", drop_params=False) assert "instructions" not in result, "Instructions should be dropped" assert result.get("temperature") == 0.7, "Other params should be preserved" @@ -94,25 +80,15 @@ class TestXAIResponsesAPITransformation: # Test with default XAI API base url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.x.ai/v1/responses" - ), f"Expected XAI responses endpoint, got {url}" + assert url == "https://api.x.ai/v1/responses", f"Expected XAI responses endpoint, got {url}" # Test with custom api_base - custom_url = config.get_complete_url( - api_base="https://custom.x.ai/v1", litellm_params={} - ) - assert ( - custom_url == "https://custom.x.ai/v1/responses" - ), f"Expected custom endpoint, got {custom_url}" + custom_url = config.get_complete_url(api_base="https://custom.x.ai/v1", litellm_params={}) + assert custom_url == "https://custom.x.ai/v1/responses", f"Expected custom endpoint, got {custom_url}" # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.x.ai/v1/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.x.ai/v1/responses" - ), "Should handle trailing slash" + url_with_slash = config.get_complete_url(api_base="https://api.x.ai/v1/", litellm_params={}) + assert url_with_slash == "https://api.x.ai/v1/responses", "Should handle trailing slash" def test_web_search_tool_transformation(self): """Test that web_search tools are transformed to XAI format""" @@ -173,9 +149,7 @@ class TestXAIResponsesAPITransformation: config = XAIResponsesAPIConfig() params = ResponsesAPIOptionalRequestParams( - tools=[ - {"type": "web_search", "excluded_domains": ["example.com", "test.com"]} - ] + tools=[{"type": "web_search", "excluded_domains": ["example.com", "test.com"]}] ) result = config.map_openai_params( @@ -338,9 +312,7 @@ class TestXAIResponsesToolUsageAttach: def test_server_side_tool_usage_details_from_usage_attr(self): usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) setattr(usage, "server_side_tool_usage_details", self._TOOL_DETAILS) - details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( - usage - ) + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage(usage) assert details == self._TOOL_DETAILS def test_server_side_tool_usage_details_from_model_extra(self): @@ -350,16 +322,11 @@ class TestXAIResponsesToolUsageAttach: total_tokens=15, server_side_tool_usage_details=self._TOOL_DETAILS, ) - details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( - usage - ) + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage(usage) assert details == self._TOOL_DETAILS def test_server_side_tool_usage_details_from_usage_none(self): - assert ( - XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage(None) - is None - ) + assert XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage(None) is None assert ( XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1) @@ -368,17 +335,13 @@ class TestXAIResponsesToolUsageAttach: ) def test_attach_noop_when_usage_missing(self): - response = ResponsesAPIResponse.model_construct( - id="resp_1", created_at=0, output=[], usage=None - ) + response = ResponsesAPIResponse.model_construct(id="resp_1", created_at=0, output=[], usage=None) XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) assert response.usage is None def test_attach_noop_when_details_missing(self): usage = ResponseAPIUsage(input_tokens=3, output_tokens=1, total_tokens=4) - response = ResponsesAPIResponse.model_construct( - id="resp_2", created_at=0, output=[], usage=usage - ) + response = ResponsesAPIResponse.model_construct(id="resp_2", created_at=0, output=[], usage=usage) XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) assert isinstance(response.usage, ResponseAPIUsage) @@ -389,32 +352,53 @@ class TestXAIResponsesToolUsageAttach: total_tokens=120, server_side_tool_usage_details=self._TOOL_DETAILS, ) - response = ResponsesAPIResponse.model_construct( - id="resp_3", created_at=0, output=[], usage=usage - ) + response = ResponsesAPIResponse.model_construct(id="resp_3", created_at=0, output=[], usage=usage) XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) assert isinstance(response.usage, Usage) assert response.usage.prompt_tokens == 100 assert response.usage.completion_tokens == 20 - assert getattr(response.usage, "server_side_tool_usage_details") == ( - self._TOOL_DETAILS - ) + assert getattr(response.usage, "server_side_tool_usage_details") == (self._TOOL_DETAILS) assert response.usage.prompt_tokens_details is not None assert response.usage.prompt_tokens_details.web_search_requests == 2 def test_attach_updates_existing_chat_usage_in_place(self): usage = Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10) setattr(usage, "server_side_tool_usage_details", self._TOOL_DETAILS) - response = ResponsesAPIResponse.model_construct( - id="resp_4", created_at=0, output=[], usage=usage - ) + response = ResponsesAPIResponse.model_construct(id="resp_4", created_at=0, output=[], usage=usage) XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) assert response.usage is usage assert usage.prompt_tokens_details is not None assert usage.prompt_tokens_details.web_search_requests == 2 + def test_chat_bridge_retransform_after_attach_keeps_tool_usage(self): + """completion(..., web_search_options={}) re-converts usage after xAI attach.""" + from litellm.responses.utils import ResponseAPILoggingUtils + + usage = ResponseAPIUsage( + input_tokens=100, + output_tokens=20, + total_tokens=120, + server_side_tool_usage_details=self._TOOL_DETAILS, + ) + response = ResponsesAPIResponse.model_construct(id="resp_bridge", created_at=0, output=[], usage=usage) + XAIResponsesAPIConfig._attach_server_side_tool_usage_details_to_usage(response) + assert isinstance(response.usage, Usage) + + bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response.usage) + assert bridged is response.usage + assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS + assert bridged.prompt_tokens_details is not None + assert bridged.prompt_tokens_details.web_search_requests == 2 + + from_dump = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(bridged.model_dump()) + assert from_dump.prompt_tokens == 100 + assert from_dump.completion_tokens == 20 + assert getattr(from_dump, "server_side_tool_usage_details") == self._TOOL_DETAILS + assert from_dump.prompt_tokens_details is not None + assert from_dump.prompt_tokens_details.web_search_requests == 2 + def test_transform_streaming_response_completed_attaches_tool_usage(self): config = XAIResponsesAPIConfig() chunk = { @@ -431,15 +415,11 @@ class TestXAIResponsesToolUsageAttach: }, }, } - event = config.transform_streaming_response( - model="grok-4.3", parsed_chunk=chunk, logging_obj=MagicMock() - ) + event = config.transform_streaming_response(model="grok-4.3", parsed_chunk=chunk, logging_obj=MagicMock()) assert isinstance(event, ResponseCompletedEvent) assert isinstance(event.response.usage, Usage) - assert getattr(event.response.usage, "server_side_tool_usage_details") == ( - self._TOOL_DETAILS - ) + assert getattr(event.response.usage, "server_side_tool_usage_details") == (self._TOOL_DETAILS) assert event.response.usage.prompt_tokens_details is not None assert event.response.usage.prompt_tokens_details.web_search_requests == 2 @@ -452,9 +432,7 @@ class TestXAIResponsesToolUsageAttach: "content_index": 0, "delta": "hi", } - event = config.transform_streaming_response( - model="grok-4.3", parsed_chunk=chunk, logging_obj=MagicMock() - ) + event = config.transform_streaming_response(model="grok-4.3", parsed_chunk=chunk, logging_obj=MagicMock()) assert getattr(event, "type", None) is not None assert not isinstance( event, diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 0141cf5d96a..7ff690ade13 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -7,9 +7,7 @@ from unittest.mock import MagicMock, patch import pytest from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -54,9 +52,7 @@ class TestResponsesAPIRequestUtils: # Setup model = "gpt-4o" config = OpenAIResponsesAPIConfig() - optional_params = ResponsesAPIOptionalRequestParams( - {"temperature": 0.7, "unsupported_param": "value"} - ) + optional_params = ResponsesAPIOptionalRequestParams({"temperature": 0.7, "unsupported_param": "value"}) # Execute and Assert with pytest.raises(litellm.UnsupportedParamsError) as excinfo: @@ -90,9 +86,7 @@ class TestResponsesAPIRequestUtils: assert result == {"temperature": 0.7} @pytest.mark.parametrize("request_drop_params", [None, False]) - def test_get_optional_params_responses_api_still_raises_without_drop( - self, monkeypatch, request_drop_params - ): + def test_get_optional_params_responses_api_still_raises_without_drop(self, monkeypatch, request_drop_params): """Absent or False request-level drop_params must not suppress the unsupported-param error""" monkeypatch.setattr(litellm, "drop_params", False) config = OpenAIResponsesAPIConfig() @@ -119,9 +113,7 @@ class TestResponsesAPIRequestUtils: } # Execute - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) # Assert assert "temperature" in result @@ -147,40 +139,31 @@ class TestResponsesAPIRequestUtils: ) # Execute - result = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id( - encoded_id - ) + result = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(encoded_id) # Assert assert result == original_response_id # Test with a non-encoded ID plain_id = "resp_xyz789" - result_plain = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id( - plain_id - ) + result_plain = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(plain_id) assert result_plain == plain_id def test_update_responses_api_response_id_with_model_id_handles_dict(self): """Ensure _update_responses_api_response_id_with_model_id works with dict input""" responses_api_response = {"id": "resp_abc123"} litellm_metadata = {"model_info": {"id": "gpt-4o"}} - updated = ( - ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=responses_api_response, - custom_llm_provider="openai", - litellm_metadata=litellm_metadata, - ) + updated = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=responses_api_response, + custom_llm_provider="openai", + litellm_metadata=litellm_metadata, ) assert updated["id"] != "resp_abc123" - decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id( - updated["id"] - ) + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(updated["id"]) assert decoded.get("response_id") == "resp_abc123" assert decoded.get("model_id") == "gpt-4o" assert decoded.get("custom_llm_provider") == "openai" - def test_update_responses_api_response_id_with_model_id_is_idempotent_for_litellm_ids(self): raw = "resp_" + "a" * 48 litellm_metadata = {"model_info": {"id": "model-123"}} @@ -207,9 +190,7 @@ class TestResponsesAPIRequestUtils: model_id=None, container_id="cntr_upstream_abc", ) - assert "None" not in base64.b64decode( - encoded.replace("cntr_", "").encode("utf-8") - ).decode("utf-8") + assert "None" not in base64.b64decode(encoded.replace("cntr_", "").encode("utf-8")).decode("utf-8") decoded = ResponsesAPIRequestUtils._decode_container_id(encoded) assert decoded.get("custom_llm_provider") == "azure" assert decoded.get("model_id") is None @@ -217,12 +198,8 @@ class TestResponsesAPIRequestUtils: def test_decode_container_id_legacy_literal_none_model_id(self): """IDs encoded before the None fix should decode without a bogus model_id.""" - legacy_inner = ( - "litellm:custom_llm_provider:azure;model_id:None;container_id:cntr_x" - ) - legacy_id = "cntr_" + base64.b64encode(legacy_inner.encode("utf-8")).decode( - "utf-8" - ) + legacy_inner = "litellm:custom_llm_provider:azure;model_id:None;container_id:cntr_x" + legacy_id = "cntr_" + base64.b64encode(legacy_inner.encode("utf-8")).decode("utf-8") decoded = ResponsesAPIRequestUtils._decode_container_id(legacy_id) assert decoded.get("model_id") is None assert decoded.get("custom_llm_provider") == "azure" @@ -264,19 +241,14 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert assert isinstance(result, Usage) assert result.prompt_tokens == 10 assert result.completion_tokens == 20 assert result.total_tokens == 30 - assert ( - result.prompt_tokens_details - and result.prompt_tokens_details.cached_tokens == 2 - ) + assert result.prompt_tokens_details and result.prompt_tokens_details.cached_tokens == 2 def test_transform_response_api_usage_with_none_values(self): """Test transformation handles None values properly""" @@ -289,9 +261,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert assert result.prompt_tokens == 0 @@ -310,9 +280,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert assert result.prompt_tokens == 15 @@ -349,9 +317,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert - verify basic token counts assert isinstance(result, Usage) @@ -386,9 +352,7 @@ class TestResponseAPILoggingUtils: }, } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert result.prompt_tokens_details is not None assert result.prompt_tokens_details.cache_write_tokens == 10059 @@ -417,9 +381,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert - all token detail types should be preserved assert result.prompt_tokens_details is not None @@ -451,9 +413,7 @@ class TestResponseAPILoggingUtils: }, } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert result.prompt_tokens_details is not None assert result.prompt_tokens_details.text_tokens == 8 @@ -475,9 +435,7 @@ class TestResponseAPILoggingUtils: "output_token_details": {"text_tokens": 2, "audio_tokens": 98}, } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert result.prompt_tokens_details is not None assert result.prompt_tokens_details.text_tokens == 10 @@ -487,6 +445,57 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.text_tokens == 20 assert result.completion_tokens_details.audio_tokens is None + def test_transform_already_chat_usage_passthrough_keeps_tool_details(self): + """xAI Responses converts usage to chat Usage before the chat bridge re-runs this helper.""" + details = {"web_search_calls": 2, "x_search_calls": 0} + usage = Usage( + prompt_tokens=100, + completion_tokens=20, + total_tokens=120, + prompt_tokens_details={"web_search_requests": 2}, + ) + setattr(usage, "server_side_tool_usage_details", details) + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result is usage + assert getattr(result, "server_side_tool_usage_details") == details + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.web_search_requests == 2 + + def test_transform_chat_shaped_usage_dict_keeps_tool_details(self): + """Streaming chat bridge dumps already-converted Usage as a prompt_tokens dict.""" + details = { + "web_search_calls": 3, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + "image_generation_calls": 0, + } + usage = { + "prompt_tokens": 50, + "completion_tokens": 10, + "total_tokens": 60, + "prompt_tokens_details": {"web_search_requests": 3, "cached_tokens": 8}, + "completion_tokens_details": {"reasoning_tokens": 4}, + "server_side_tool_usage_details": details, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert isinstance(result, Usage) + assert result.prompt_tokens == 50 + assert result.completion_tokens == 10 + assert result.total_tokens == 60 + assert getattr(result, "server_side_tool_usage_details") == details + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.web_search_requests == 3 + assert result.prompt_tokens_details.cached_tokens == 8 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.reasoning_tokens == 4 + class TestResponsesAPIProviderSpecificParams: """ @@ -503,9 +512,7 @@ class TestResponsesAPIProviderSpecificParams: } # Should not raise any exception - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) assert "temperature" in result def test_provider_specific_params_no_crash_with_openai(self): @@ -517,9 +524,7 @@ class TestResponsesAPIProviderSpecificParams: } # Should not raise any exception - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) assert "temperature" in result def test_provider_specific_params_no_crash_with_vertex_ai(self): @@ -531,9 +536,7 @@ class TestResponsesAPIProviderSpecificParams: } # Should not raise any exception - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) assert "temperature" in result From 4a536098e129135a78949c43a8e9263b725ab0af Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 8 Aug 2026 16:58:25 -0700 Subject: [PATCH 033/245] style: ruff format xAI cost calculator and responses transform --- litellm/llms/xai/cost_calculator.py | 4 +--- litellm/llms/xai/responses/transformation.py | 12 +++--------- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index ebeac61fb43..a32e2af0959 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -17,9 +17,7 @@ if TYPE_CHECKING: _DEFAULT_WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 -def apply_server_side_tool_usage_details_to_usage( - usage: Usage, details: Mapping[str, Any] | None -) -> None: +def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping[str, Any] | None) -> None: """ Attach server_side_tool_usage_details and mirror web_search_calls onto prompt_tokens_details.web_search_requests for built-in tool cost gating. diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index de041854b3a..35be5602029 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -96,9 +96,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): attaching server_side_tool_usage_details here, stream=true web_search usage is dropped when usage is normalized for billing. """ - event = super().transform_streaming_response( - model=model, parsed_chunk=parsed_chunk, logging_obj=logging_obj - ) + event = super().transform_streaming_response(model=model, parsed_chunk=parsed_chunk, logging_obj=logging_obj) if isinstance( event, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent), @@ -117,9 +115,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): details = getattr(usage, "server_side_tool_usage_details", None) if details is not None: return details - model_extra = getattr(usage, "model_extra", None) or getattr( - usage, "__pydantic_extra__", None - ) + model_extra = getattr(usage, "model_extra", None) or getattr(usage, "__pydantic_extra__", None) if isinstance(model_extra, dict): return model_extra.get("server_side_tool_usage_details") return None @@ -131,9 +127,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): if response.usage is None: return - details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage( - response.usage - ) + details = XAIResponsesAPIConfig._server_side_tool_usage_details_from_usage(response.usage) if details is None: return From 21f742041a6cede248d57efc53c2dd1d969330e2 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Sat, 8 Aug 2026 20:05:43 -0700 Subject: [PATCH 034/245] fix(xai): type tool usage details helpers without Any --- litellm/llms/xai/cost_calculator.py | 4 +-- litellm/llms/xai/responses/transformation.py | 26 ++++++++++++-------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index a32e2af0959..0f117c586f1 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -5,7 +5,7 @@ Helper util for handling XAI-specific cost calculation """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import PromptTokensDetailsWrapper, Usage @@ -17,7 +17,7 @@ if TYPE_CHECKING: _DEFAULT_WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 -def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping[str, Any] | None) -> None: +def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping[str, object] | None) -> None: """ Attach server_side_tool_usage_details and mirror web_search_calls onto prompt_tokens_details.web_search_requests for built-in tool cost gating. diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 35be5602029..3dc11cd2d92 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -14,6 +15,7 @@ from litellm.llms.xai.cost_calculator import ( from litellm.responses.utils import ResponseAPILoggingUtils from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( + ResponseAPIUsage, ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent, @@ -107,18 +109,22 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return event @staticmethod - def _server_side_tool_usage_details_from_usage(usage: Any) -> Any: + def _server_side_tool_usage_details_from_usage( + usage: Usage | ResponseAPIUsage | Mapping[str, object] | None, + ) -> Mapping[str, object] | None: if usage is None: return None - if isinstance(usage, dict): - return usage.get("server_side_tool_usage_details") - details = getattr(usage, "server_side_tool_usage_details", None) - if details is not None: - return details - model_extra = getattr(usage, "model_extra", None) or getattr(usage, "__pydantic_extra__", None) - if isinstance(model_extra, dict): - return model_extra.get("server_side_tool_usage_details") - return None + if isinstance(usage, Mapping): + details = usage.get("server_side_tool_usage_details") + else: + details = getattr(usage, "server_side_tool_usage_details", None) + if details is None: + model_extra = getattr(usage, "model_extra", None) or getattr(usage, "__pydantic_extra__", None) + if isinstance(model_extra, Mapping): + details = model_extra.get("server_side_tool_usage_details") + if not isinstance(details, Mapping): + return None + return details @staticmethod def _attach_server_side_tool_usage_details_to_usage( From 30c4898de9cea90e777a43a5260fe77011acdb5b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 8 Aug 2026 20:16:22 -0700 Subject: [PATCH 035/245] fix(ui): hide admin-only Logs tabs from roles that cannot call their endpoints The Logs nav entry is open to internal users so they can read their own request logs, but the page rendered all four tabs unconditionally. Audit Logs calls GET /audit and Deleted Teams calls GET /v2/team/list?status=deleted, neither of which an internal user is permitted to call, so the page fired requests that came back 401. Gate both tabs on new viewAuditLogs / viewDeletedTeams capabilities, using the same CAPABILITY_ROLES map and useCan hook introduced for Tool Policies. Hiding a tab drops its panel from the tree entirely, so the request is never issued rather than issued and rejected. Selecting a tab also mapped index 0 to "request logs" and every other index to "audit logs", which activated the audit panel whenever a user opened Deleted Keys or Deleted Teams. Derive the active tab from the visible tab list instead, so the mapping survives tabs being filtered out. --- .../view_logs/index.integration.test.tsx | 104 ++++++++++++++++++ .../src/components/view_logs/index.test.tsx | 79 ++++++++++++- .../src/components/view_logs/index.tsx | 91 +++++++++------ .../src/utils/capabilities.test.ts | 13 +++ .../src/utils/capabilities.ts | 2 + 5 files changed, 254 insertions(+), 35 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/view_logs/index.integration.test.tsx diff --git a/ui/litellm-dashboard/src/components/view_logs/index.integration.test.tsx b/ui/litellm-dashboard/src/components/view_logs/index.integration.test.tsx new file mode 100644 index 00000000000..b86ad015b91 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/index.integration.test.tsx @@ -0,0 +1,104 @@ +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import SpendLogsTable from "./index"; +import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; + +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); + +vi.mock("./RequestLogsPanel", () => ({ + default: function RequestLogsPanelMock() { + return
; + }, +})); + +const fetchMock = vi.fn(); + +const jsonResponse = (body: unknown) => ({ + ok: true, + status: 200, + statusText: "OK", + json: async () => body, +}); + +const requestedUrls = () => fetchMock.mock.calls.map(([url]) => String(url)); + +const emptyAuditLogs = { audit_logs: [], total: 0, page: 1, page_size: 50, total_pages: 0 }; + +const defaultProps = { + accessToken: "sk-test", + token: "jwt-test", + userRole: "Admin", + userID: "user-1", + premiumUser: true, +}; + +const renderAs = (sessionRole: string) => { + useAuthorizedMock.mockReturnValue({ accessToken: "sk-test", userRole: sessionRole, premiumUser: true }); + return renderWithProviders(); +}; + +describe("SpendLogsTable network access by role", () => { + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + fetchMock.mockImplementation(async (url: string) => { + if (String(url).includes("/audit")) { + return jsonResponse(emptyAuditLogs); + } + if (String(url).includes("/v2/team/list")) { + return jsonResponse({ teams: [] }); + } + return jsonResponse({ keys: [], total_count: 0 }); + }); + vi.stubGlobal("fetch", fetchMock); + }); + + it("fires neither the audit nor the deleted-teams request for an internal user", async () => { + const user = userEvent.setup(); + renderAs("Internal User"); + + // Liveness gate: the sibling Deleted Keys panel does reach the network, so a + // silent absence below means the gate worked, not that nothing rendered. + await waitFor(() => expect(requestedUrls().some((url) => url.includes("/key/list"))).toBe(true)); + + await user.click(screen.getByRole("tab", { name: "Deleted Keys" })); + await user.click(screen.getByRole("tab", { name: "Request Logs" })); + + expect(requestedUrls().filter((url) => url.includes("/audit"))).toEqual([]); + expect(requestedUrls().filter((url) => url.includes("/v2/team/list"))).toEqual([]); + }); + + it("fetches deleted teams and audit logs for an admin", async () => { + const user = userEvent.setup(); + renderAs("Admin"); + + await waitFor(() => + expect(requestedUrls().some((url) => url.includes("/v2/team/list") && url.includes("status=deleted"))).toBe(true), + ); + + expect(requestedUrls().filter((url) => url.includes("/audit"))).toEqual([]); + + await user.click(screen.getByRole("tab", { name: "Audit Logs" })); + + await waitFor(() => expect(requestedUrls().some((url) => url.includes("/audit"))).toBe(true)); + }); + + it("leaves the audit request unsent when an admin selects a tab after Audit Logs", async () => { + const user = userEvent.setup(); + renderAs("Admin"); + + await user.click(screen.getByRole("tab", { name: "Deleted Teams" })); + + expect(screen.getByRole("tab", { name: "Deleted Teams" })).toHaveAttribute("aria-selected", "true"); + expect(requestedUrls().filter((url) => url.includes("/audit"))).toEqual([]); + + await user.click(screen.getByRole("tab", { name: "Audit Logs" })); + + await waitFor(() => expect(requestedUrls().some((url) => url.includes("/audit"))).toBe(true)); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx index b2e77ec7fd5..785fa0cc6f8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx @@ -1,9 +1,15 @@ import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import SpendLogsTable from "./index"; import { renderWithProviders } from "../../../tests/test-utils"; +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); + vi.mock("./RequestLogsPanel", () => ({ default: function RequestLogsPanelMock({ isActive }: { isActive: boolean }) { return
{isActive ? "active" : "inactive"}
; @@ -36,9 +42,18 @@ const defaultProps = { premiumUser: false, }; +const renderAs = (sessionRole: string) => { + useAuthorizedMock.mockReturnValue({ userRole: sessionRole }); + return renderWithProviders(); +}; + describe("SpendLogsTable", () => { + beforeEach(() => { + useAuthorizedMock.mockReturnValue({ userRole: "Admin" }); + }); + it("renders the four log tabs", () => { - renderWithProviders(); + renderAs("Admin"); for (const label of ["Request Logs", "Audit Logs", "Deleted Keys", "Deleted Teams"]) { expect(screen.getByRole("tab", { name: label })).toBeInTheDocument(); @@ -47,7 +62,7 @@ describe("SpendLogsTable", () => { it("marks only the visible tab's panel active so background tabs do not query", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderAs("Admin"); expect(screen.getByTestId("request-logs-panel")).toHaveTextContent("active"); @@ -57,8 +72,64 @@ describe("SpendLogsTable", () => { expect(screen.getByTestId("request-logs-panel")).toHaveTextContent("inactive"); }); + describe("admin-only tabs", () => { + it.each(["Internal User", "Internal Viewer"])("hides Audit Logs and Deleted Teams from %s", (role) => { + renderAs(role); + + expect(screen.getByRole("tab", { name: "Request Logs" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Deleted Keys" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Audit Logs" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Deleted Teams" })).not.toBeInTheDocument(); + }); + + it("never mounts the panels that call the admin-only endpoints for an internal user", () => { + renderAs("Internal User"); + + expect(screen.queryByTestId("audit-logs-panel")).not.toBeInTheDocument(); + expect(screen.queryByTestId("deleted-teams-page")).not.toBeInTheDocument(); + expect(screen.getByTestId("deleted-keys-page")).toBeInTheDocument(); + }); + }); + + describe("tab index mapping", () => { + it("activates the panel the admin selected, not the one at the old hardcoded index", async () => { + const user = userEvent.setup(); + renderAs("Admin"); + + await user.click(screen.getByRole("tab", { name: "Deleted Keys" })); + + expect(screen.getByTestId("audit-logs-panel")).toHaveTextContent("inactive"); + expect(screen.getByTestId("request-logs-panel")).toHaveTextContent("inactive"); + }); + + it("keeps the audit panel inert when an admin selects the last tab", async () => { + const user = userEvent.setup(); + renderAs("Admin"); + + await user.click(screen.getByRole("tab", { name: "Deleted Teams" })); + + expect(screen.getByTestId("audit-logs-panel")).toHaveTextContent("inactive"); + expect(screen.getByTestId("deleted-teams-page")).toBeInTheDocument(); + }); + + it("selects the last visible tab for an internal user and returns to Request Logs", async () => { + const user = userEvent.setup(); + renderAs("Internal User"); + + await user.click(screen.getByRole("tab", { name: "Deleted Keys" })); + + expect(screen.getByTestId("deleted-keys-page")).toBeInTheDocument(); + expect(screen.getByTestId("request-logs-panel")).toHaveTextContent("inactive"); + + await user.click(screen.getByRole("tab", { name: "Request Logs" })); + + expect(screen.getByTestId("request-logs-panel")).toHaveTextContent("active"); + }); + }); + describe("auth-not-ready guard", () => { it("shows a loading spinner when credentials are not yet resolved", () => { + useAuthorizedMock.mockReturnValue({ userRole: "Admin" }); renderWithProviders(); expect(document.querySelector(".ant-spin")).toBeInTheDocument(); @@ -66,7 +137,7 @@ describe("SpendLogsTable", () => { }); it("renders the tabs (no spinner) once all credentials are present", () => { - renderWithProviders(); + renderAs("Admin"); expect(document.querySelector(".ant-spin")).not.toBeInTheDocument(); expect(screen.getByRole("tab", { name: "Request Logs" })).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 8e7423e3fae..7269564dcec 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -1,5 +1,6 @@ import { useState } from "react"; import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import DeletedKeysPage from "../DeletedKeysPage/DeletedKeysPage"; import DeletedTeamsPage from "../DeletedTeamsPage/DeletedTeamsPage"; import AuditLogsPanel from "./AuditLogsPanel"; @@ -14,8 +15,22 @@ interface SpendLogsTableProps { premiumUser: boolean; } +type LogsTabId = "request logs" | "audit logs" | "deleted keys" | "deleted teams"; + +interface LogsTab { + id: LogsTabId; + label: string; +} + +const REQUEST_LOGS_TAB: LogsTab = { id: "request logs", label: "Request Logs" }; +const AUDIT_LOGS_TAB: LogsTab = { id: "audit logs", label: "Audit Logs" }; +const DELETED_KEYS_TAB: LogsTab = { id: "deleted keys", label: "Deleted Keys" }; +const DELETED_TEAMS_TAB: LogsTab = { id: "deleted teams", label: "Deleted Teams" }; + export default function SpendLogsTable({ accessToken, token, userRole, userID, premiumUser }: SpendLogsTableProps) { - const [activeTab, setActiveTab] = useState("request logs"); + const [activeTab, setActiveTab] = useState(REQUEST_LOGS_TAB.id); + const canViewAuditLogs = useCan("viewAuditLogs"); + const canViewDeletedTeams = useCan("viewDeletedTeams"); if (!accessToken || !token || !userRole || !userID) { return ( @@ -25,41 +40,55 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p ); } + const tabs: LogsTab[] = [ + REQUEST_LOGS_TAB, + ...(canViewAuditLogs ? [AUDIT_LOGS_TAB] : []), + DELETED_KEYS_TAB, + ...(canViewDeletedTeams ? [DELETED_TEAMS_TAB] : []), + ]; + + const renderPanel = (tabId: LogsTabId) => { + switch (tabId) { + case "request logs": + return ( + + ); + case "audit logs": + return ( + + ); + case "deleted keys": + return ; + case "deleted teams": + return ; + } + }; + return (
- setActiveTab(index === 0 ? "request logs" : "audit logs")}> + setActiveTab(tabs[index].id)}> - Request Logs - Audit Logs - Deleted Keys - Deleted Teams + {tabs.map((tab) => ( + {tab.label} + ))} - - - - - - - - - - - - + {tabs.map((tab) => ( + {renderPanel(tab.id)} + ))}
diff --git a/ui/litellm-dashboard/src/utils/capabilities.test.ts b/ui/litellm-dashboard/src/utils/capabilities.test.ts index f48609b0b9d..611c9626065 100644 --- a/ui/litellm-dashboard/src/utils/capabilities.test.ts +++ b/ui/litellm-dashboard/src/utils/capabilities.test.ts @@ -18,6 +18,19 @@ describe("hasCapability", () => { ); }); +describe.each(["viewAuditLogs", "viewDeletedTeams"] as const)("hasCapability - %s", (capability) => { + it.each(["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer"])("should grant it to %s", (role) => { + expect(hasCapability(role, capability)).toBe(true); + }); + + it.each(["Internal User", "Internal Viewer", "App User", "Org Admin", "Unknown Role", "", null, undefined])( + "should deny it to %s", + (role) => { + expect(hasCapability(role, capability)).toBe(false); + }, + ); +}); + describe("rolesWithCapability", () => { it("should return a copy so callers cannot mutate the capability map", () => { const roles = rolesWithCapability("viewToolPolicies"); diff --git a/ui/litellm-dashboard/src/utils/capabilities.ts b/ui/litellm-dashboard/src/utils/capabilities.ts index 77ead2568fb..f0847cc3400 100644 --- a/ui/litellm-dashboard/src/utils/capabilities.ts +++ b/ui/litellm-dashboard/src/utils/capabilities.ts @@ -2,6 +2,8 @@ import { all_admin_roles } from "./roles"; const CAPABILITY_ROLES = { viewToolPolicies: all_admin_roles, + viewAuditLogs: all_admin_roles, + viewDeletedTeams: all_admin_roles, } as const satisfies Record; export type Capability = keyof typeof CAPABILITY_ROLES; From 6a540a1bf848129dc16228d1b23f92120d7a7f03 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 8 Aug 2026 20:17:12 -0700 Subject: [PATCH 036/245] fix(ui): gate organization and agent usage views behind capabilities The Usage page admits internal users because their own usage view works, but the entity breakdown selector inside it also offered Organization Usage, so picking it fired /organization/daily/activity and collected a 401. Neither that route nor /agent/daily/activity appears in any non-admin route list, so both are default-deny. The team breakdown leaked the second one too: it fetches agent activity unconditionally to fill its Top Agents card, which 401s for the same roles. Adds viewOrganizationUsage and viewAgentUsage to the existing capability map and points the selector option, the page section, and the fetch's enabled flag at the same capability, so a role that cannot call the endpoint never sees the breakdown and never issues the request. The team and tag breakdowns, which internal users can read, are untouched, and the default Usage view was already one of those. --- .../EntityUsage/EntityUsage.test.tsx | 43 ++++++++++++++++++- .../components/EntityUsage/EntityUsage.tsx | 34 +++++++++++---- .../components/UsagePageView.test.tsx | 25 +++++++++++ .../_components/components/UsagePageView.tsx | 9 ++-- .../UsageViewSelect/UsageViewSelect.test.tsx | 29 +++++++++++-- .../UsageViewSelect/UsageViewSelect.tsx | 20 +++++---- .../src/utils/capabilities.test.ts | 36 ++++++++++------ .../src/utils/capabilities.ts | 2 + 8 files changed, 160 insertions(+), 38 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index 82ca66b10c0..11528117f1e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -1,4 +1,4 @@ -import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { act, cleanup, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "@/components/networking"; import EntityUsage from "./EntityUsage"; @@ -856,6 +856,47 @@ describe("EntityUsage", () => { expect(logo.getAttribute("src")).toContain("openai_small"); }); + describe("capability gating", () => { + it.each([ + ["organization", () => mockOrganizationDailyActivityCall, "Organization Spend Overview"], + ["agent", () => mockAgentDailyActivityCall, "Agent Spend Overview"], + ] as const)("fetches %s activity for an admin but not for an internal user", async (entityType, call, heading) => { + render(); + await waitFor(() => { + expect(call()).toHaveBeenCalled(); + }); + + cleanup(); + call().mockClear(); + + render(); + expect(await screen.findByText(heading)).toBeInTheDocument(); + expect(call()).not.toHaveBeenCalled(); + }); + + it("keeps the team breakdown but drops its agent sub-fetch for an internal user", async () => { + render(); + + await waitFor(() => { + expect(mockTeamDailyActivityCall).toHaveBeenCalled(); + }); + expect(screen.getByText("Team Spend Overview")).toBeInTheDocument(); + + expect(mockAgentDailyActivityCall).not.toHaveBeenCalled(); + expect(screen.queryByText("Agent Activity")).not.toBeInTheDocument(); + expect(screen.queryByText("Top Agents Driving Spend")).not.toBeInTheDocument(); + }); + + it("keeps the tag breakdown for an internal user", async () => { + render(); + + await waitFor(() => { + expect(mockTagDailyActivityCall).toHaveBeenCalled(); + }); + expect(screen.getByText("Tag Spend Overview")).toBeInTheDocument(); + }); + }); + it("renders a letter avatar instead of an img for an unknown provider slug", async () => { const spendDataUnknownProvider = { ...mockSpendData, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 4d44791d1a9..5a0f2abf15b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -2,6 +2,7 @@ import useTeams from "@/app/(dashboard)/hooks/useTeams"; import { BarChart, DonutChart } from "@/components/shared/charts"; import { MoneyCell } from "@/components/shared/table_cells"; import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { hasCapability, type Capability } from "@/utils/capabilities"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { Card, @@ -108,7 +109,19 @@ const ENTITY_FETCH_FNS: Record Promise> = { user: userDailyActivityCall, }; -const EntityUsage: React.FC = ({ accessToken, entityType, entityId, entityList, dateValue }) => { +const ENTITY_CAPABILITIES: Partial> = { + organization: "viewOrganizationUsage", + agent: "viewAgentUsage", +}; + +const EntityUsage: React.FC = ({ + accessToken, + entityType, + entityId, + entityList, + userRole, + dateValue, +}) => { const { teams } = useTeams(); const [selectedTags, setSelectedTags] = useState([]); const [modelViewType, setModelViewType] = useState("groups"); @@ -125,7 +138,11 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti }, [entityType, selectedTags]); const fetchFn = ENTITY_FETCH_FNS[entityType]; - const enabled = !!accessToken && !!startTime && !!endTime; + const entityCapability = ENTITY_CAPABILITIES[entityType]; + const canViewEntity = entityCapability === undefined || hasCapability(userRole, entityCapability); + const showAgentBreakdown = entityType === "team" && hasCapability(userRole, "viewAgentUsage"); + const hasRequestWindow = !!accessToken && !!startTime && !!endTime; + const enabled = hasRequestWindow && canViewEntity; const { data: spendDataRaw, @@ -150,7 +167,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti } = usePaginatedDailyActivity({ fetchFn: agentDailyActivityCall, args: [accessToken, startTime, endTime, null], - enabled: enabled && entityType === "team", + enabled: enabled && showAgentBreakdown, }); const agentSpendData = agentSpendDataRaw as unknown as EntitySpendData; @@ -158,7 +175,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti const modelBreakdownKey = modelViewType === "groups" ? "model_groups" : "models"; const modelMetrics = processActivityData(spendData, modelBreakdownKey, teams || []); const keyMetrics = processActivityData(spendData, "api_keys", teams || []); - const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {}; + const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {}; const getTopModels = () => { const modelSpend: { [key: string]: any } = {}; @@ -621,8 +638,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti - {/* Top Agents - only for team entity type */} - {entityType === "team" && ( + {showAgentBreakdown && ( Top Agents Driving Spend @@ -708,7 +724,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti ), }, - ...(entityType === "team" + ...(showAgentBreakdown ? [{ key: "agents", label: "Agent Activity", content: }] : []), { @@ -757,7 +773,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti } /> )} - {agentIsFetchingMore && entityType === "team" && ( + {agentIsFetchingMore && showAgentBreakdown && ( = ({ accessToken, entityType, enti } /> )} - {agentCancelled && entityType === "team" && ( + {agentCancelled && showAgentBreakdown && ( { userId: "user-123", userEmail: "test@example.com", userRole: "Internal User", + userRoleLabel: "Internal User", + isViewOnly: false, premiumUser: true, disabledPersonalKeyCreation: false, showSSOBanner: false, @@ -861,6 +863,29 @@ describe("UsagePage", () => { }); }); + // The select hides both views from a non-admin, so this drives the section + // gate directly through the mocked select, which always offers every option. + it.each(["organization", "agent"])("should not render the %s usage view for an internal user", async (usageView) => { + mockUseAuthorized.mockReturnValue(nonAdminSession); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + const usageSelect = screen.getByTestId("usage-view-select"); + act(() => { + fireEvent.change(usageSelect, { target: { value: "team" } }); + }); + expect(screen.getAllByText("Entity Usage").length).toBeGreaterThan(0); + + act(() => { + fireEvent.change(usageSelect, { target: { value: usageView } }); + }); + expect(screen.queryByText("Entity Usage")).not.toBeInTheDocument(); + }); + describe("admin user selector", () => { it("should render user selector for admin users in global view", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index c3645d6371e..494df313ac0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -33,6 +33,7 @@ import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; +import { hasCapability } from "@/utils/capabilities"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { all_admin_roles, internalUserRoles } from "@/utils/roles"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; @@ -109,6 +110,8 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const { data: currentUser } = useCurrentUser(); const isAdmin = all_admin_roles.includes(userRole || ""); const canViewTagUsage = isAdmin || internalUserRoles.includes(userRole || ""); + const canViewOrganizationUsage = hasCapability(userRole, "viewOrganizationUsage"); + const canViewAgentUsage = hasCapability(userRole, "viewAgentUsage"); // Debounced search for user selector const [userSearchInput, setUserSearchInput] = useState(""); @@ -513,7 +516,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { setUsageView(value)} - isAdmin={isAdmin} + userRole={userRole} canViewTagUsage={canViewTagUsage} /> @@ -950,7 +953,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { )} {/* Organization Usage Panel */} - {usageView === "organization" && ( + {usageView === "organization" && canViewOrganizationUsage && ( = ({ teams, organizations }) => { /> )} - {usageView === "agent" && ( + {usageView === "agent" && canViewAgentUsage && ( { }); it("should render", () => { - render(); + render(); expect(screen.getByText("Usage View")).toBeInTheDocument(); expect(screen.getByText("Select the usage data you want to view")).toBeInTheDocument(); expect(screen.getByRole("combobox")).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "Your Usage" })).toBeInTheDocument(); }); it("should call onChange when value changes", () => { - render(); + render(); const select = screen.getByRole("combobox"); act(() => { @@ -109,14 +110,34 @@ describe("UsageViewSelect", () => { }); it("should show Tag Usage for non-admin users with tag usage permission", () => { - render(); + render(); expect(screen.getByRole("option", { name: "Tag Usage" })).toBeInTheDocument(); }); it("should hide Tag Usage for non-admin users without tag usage permission", () => { - render(); + render(); expect(screen.queryByRole("option", { name: "Tag Usage" })).not.toBeInTheDocument(); }); + + it.each(["Organization Usage", "Agent Usage (A2A)"])("should show %s to an admin", (optionName) => { + render(); + + expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument(); + }); + + // Neither /organization/daily/activity nor /agent/daily/activity admits an + // internal user, so the option that fires them must not be selectable. + it.each(["Organization Usage", "Agent Usage (A2A)"])("should hide %s from an internal user", (optionName) => { + render(); + + expect(screen.queryByRole("option", { name: optionName })).not.toBeInTheDocument(); + }); + + it.each(["Team Usage", "Tag Usage"])("should keep %s available to an internal user", (optionName) => { + render(); + + expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx index 94b483cb539..54c1d5ab7cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx @@ -11,6 +11,8 @@ import { } from "@ant-design/icons"; import { Badge, Select } from "antd"; import React from "react"; +import { hasCapability, type Capability } from "@/utils/capabilities"; +import { all_admin_roles } from "@/utils/roles"; export type UsageOption = | "global" | "my-usage" @@ -24,7 +26,7 @@ export type UsageOption = export interface UsageViewSelectProps { value: UsageOption; onChange: (value: UsageOption) => void; - isAdmin: boolean; + userRole: string | null; canViewTagUsage?: boolean; title?: string; description?: string; @@ -35,6 +37,7 @@ interface OptionConfig { label: string; description: string; icon: React.ReactNode; + capability?: Capability; adminOnly?: boolean; showForAdmin?: string; showForNonAdmin?: string; @@ -63,12 +66,9 @@ const OPTIONS: OptionConfig[] = [ { value: "organization", label: "Organization Usage", - showForAdmin: "Organization Usage", - showForNonAdmin: "Your Organization Usage", - description: "View organization-level usage", - descriptionForAdmin: "View usage across all organizations", - descriptionForNonAdmin: "View your organization's usage", + description: "View usage across all organizations", icon: , + capability: "viewOrganizationUsage", }, { value: "team", @@ -95,7 +95,7 @@ const OPTIONS: OptionConfig[] = [ label: "Agent Usage (A2A)", description: "View usage by AI agents", icon: , - adminOnly: true, + capability: "viewAgentUsage", }, { value: "user", @@ -115,14 +115,18 @@ const OPTIONS: OptionConfig[] = [ export const UsageViewSelect: React.FC = ({ value, onChange, - isAdmin, + userRole, canViewTagUsage = false, title = "Usage View", description = "Select the usage data you want to view", "data-id": dataId, }) => { + const isAdmin = all_admin_roles.includes(userRole ?? ""); const getFilteredOptions = () => { return OPTIONS.filter((option) => { + if (option.capability) { + return hasCapability(userRole, option.capability); + } if (option.value === "tag" && canViewTagUsage) { return true; } diff --git a/ui/litellm-dashboard/src/utils/capabilities.test.ts b/ui/litellm-dashboard/src/utils/capabilities.test.ts index f48609b0b9d..3f5a8ac81fb 100644 --- a/ui/litellm-dashboard/src/utils/capabilities.test.ts +++ b/ui/litellm-dashboard/src/utils/capabilities.test.ts @@ -1,21 +1,31 @@ import { describe, expect, it } from "vitest"; -import { hasCapability, rolesWithCapability } from "./capabilities"; +import { hasCapability, rolesWithCapability, type Capability } from "./capabilities"; + +const ADMIN_ROLES = ["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer"]; +const NON_ADMIN_ROLES = [ + "Internal User", + "Internal Viewer", + "App User", + "Org Admin", + "Unknown Role", + "", + null, + undefined, +]; + +const ADMIN_ONLY_CAPABILITIES: Capability[] = ["viewToolPolicies", "viewOrganizationUsage", "viewAgentUsage"]; describe("hasCapability", () => { - it.each(["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer"])( - "should grant viewToolPolicies to %s", - (role) => { - expect(hasCapability(role, "viewToolPolicies")).toBe(true); - }, - ); + describe.each(ADMIN_ONLY_CAPABILITIES)("%s", (capability) => { + it.each(ADMIN_ROLES)("should grant it to %s", (role) => { + expect(hasCapability(role, capability)).toBe(true); + }); - it.each(["Internal User", "Internal Viewer", "App User", "Org Admin", "Unknown Role", "", null, undefined])( - "should deny viewToolPolicies to %s", - (role) => { - expect(hasCapability(role, "viewToolPolicies")).toBe(false); - }, - ); + it.each(NON_ADMIN_ROLES)("should deny it to %s", (role) => { + expect(hasCapability(role, capability)).toBe(false); + }); + }); }); describe("rolesWithCapability", () => { diff --git a/ui/litellm-dashboard/src/utils/capabilities.ts b/ui/litellm-dashboard/src/utils/capabilities.ts index 77ead2568fb..c4d878b81a5 100644 --- a/ui/litellm-dashboard/src/utils/capabilities.ts +++ b/ui/litellm-dashboard/src/utils/capabilities.ts @@ -2,6 +2,8 @@ import { all_admin_roles } from "./roles"; const CAPABILITY_ROLES = { viewToolPolicies: all_admin_roles, + viewOrganizationUsage: all_admin_roles, + viewAgentUsage: all_admin_roles, } as const satisfies Record; export type Capability = keyof typeof CAPABILITY_ROLES; From 2502ee4a2ace88ac10dd8ed30e2a03fff075ce26 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 8 Aug 2026 20:34:56 -0700 Subject: [PATCH 037/245] fix(ui): gate policy and prompt lookups on an admin capability /policies/list and /prompts/list are default-deny for internal_user, but the Virtual Keys create/edit flow, the Teams forms and the Playground called them on mount, so every internal user landing on the dashboard fired two requests that 401. Add viewPolicies and viewPrompts to the capability map and use them to gate the nav entry, the form field and the fetch together, following the pattern from the Tool Policies migration. Non-admins now see no policy or prompt selector at all rather than an empty dropdown. --- .../playground/components/chat_ui/ChatUI.tsx | 56 +++---- .../components/complianceUI/ComplianceUI.tsx | 48 +++--- .../src/components/Teams.test.tsx | 55 ++++++- ui/litellm-dashboard/src/components/Teams.tsx | 68 ++++---- .../src/components/leftnav.test.tsx | 22 +++ .../src/components/leftnav.tsx | 10 +- .../organisms/create_key_button.test.tsx | 49 +++++- .../organisms/create_key_button.tsx | 145 +++++++++--------- .../policies/PolicySelector.test.tsx | 20 +++ .../components/policies/PolicySelector.tsx | 10 +- .../src/components/team/TeamInfo.test.tsx | 46 ++++++ .../src/components/team/TeamInfo.tsx | 56 +++---- .../templates/key_edit_view.test.tsx | 50 +++++- .../components/templates/key_edit_view.tsx | 87 ++++++----- .../src/utils/capabilities.test.ts | 28 ++++ .../src/utils/capabilities.ts | 2 + 16 files changed, 533 insertions(+), 219 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 57ff7906eda..0241ef8a77e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -25,6 +25,7 @@ import React, { useEffect, useRef, useState } from "react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import { v4 as uuidv4 } from "uuid"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; import PolicySelector from "@/components/policies/PolicySelector"; import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; @@ -106,6 +107,7 @@ const ChatUI: React.FC = ({ simplified = false, fixedModel, }) => { + const canViewPolicies = useCan("viewPolicies"); const [mcpServers, setMCPServers] = useState([]); const [mcpToolsets, setMCPToolsets] = useState([]); const [isToolsetsInfoModalVisible, setIsToolsetsInfoModalVisible] = useState(false); @@ -1652,32 +1654,34 @@ const ChatUI: React.FC = ({ />
-
- - Policies - - Select policy/policies to apply to this LLM API call. Policies define which guardrails are - applied based on conditions. You can set up your policies{" "} - - here - - . - - } - > - - - - -
+ {canViewPolicies && ( +
+ + Policies + + Select policy/policies to apply to this LLM API call. Policies define which guardrails are + applied based on conditions. You can set up your policies{" "} + + here + + . + + } + > + + + + +
+ )} {/* Code Interpreter Toggle - Only for Responses endpoint */} {endpointType === EndpointType.RESPONSES && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index 39346105f2a..c3b417987e6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -6,6 +6,7 @@ import { type ComplianceFramework, type CompliancePrompt, } from "@/data/compliancePrompts"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/networking"; import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; @@ -123,6 +124,7 @@ export default function ComplianceUI({ fixedModel, proxySettings, }: ComplianceUIProps) { + const canViewPolicies = useCan("viewPolicies"); const frameworks = getFrameworks(); const [policyValueToLabel, setPolicyValueToLabel] = useState>(new Map()); @@ -701,29 +703,37 @@ export default function ComplianceUI({

Test Configuration

-

Select policies, guardrails, or both to test against.

+

+ {canViewPolicies + ? "Select policies, guardrails, or both to test against." + : "Select guardrails to test against."} +

-
- - {accessToken && ( - - )} -
+ {canViewPolicies && ( + <> +
+ + {accessToken && ( + + )} +
-
-
- or -
-
+
+
+ or +
+
+ + )}