diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index ab92d334ae..c32a7b75c5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -239,6 +239,36 @@ class PrometheusLogger(CustomLogger): ), buckets=LATENCY_BUCKETS, ) + + # Request queue time metric + self.litellm_request_queue_time_metric = self._histogram_factory( + "litellm_request_queue_time_seconds", + "Time spent in request queue before processing starts (seconds)", + labelnames=self.get_labels_for_metric( + "litellm_request_queue_time_seconds" + ), + buckets=LATENCY_BUCKETS, + ) + + # Guardrail metrics + self.litellm_guardrail_latency_metric = self._histogram_factory( + "litellm_guardrail_latency_seconds", + "Latency (seconds) for guardrail execution", + labelnames=["guardrail_name", "status", "error_type", "hook_type"], + buckets=LATENCY_BUCKETS, + ) + + self.litellm_guardrail_errors_total = self._counter_factory( + "litellm_guardrail_errors_total", + "Total number of errors encountered during guardrail execution", + labelnames=["guardrail_name", "error_type", "hook_type"], + ) + + self.litellm_guardrail_requests_total = self._counter_factory( + "litellm_guardrail_requests_total", + "Total number of guardrail invocations", + labelnames=["guardrail_name", "status", "hook_type"], + ) # llm api provider budget metrics self.litellm_provider_remaining_budget_metric = self._gauge_factory( "litellm_provider_remaining_budget_metric", @@ -1262,6 +1292,22 @@ class PrometheusLogger(CustomLogger): total_time_seconds ) + # request queue time (time from arrival to processing start) + _litellm_params = kwargs.get("litellm_params", {}) or {} + queue_time_seconds = _litellm_params.get("metadata", {}).get( + "queue_time_seconds" + ) + if queue_time_seconds is not None and queue_time_seconds >= 0: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_request_queue_time_seconds" + ), + enum_values=enum_values, + ) + self.litellm_request_queue_time_metric.labels(**_labels).observe( + queue_time_seconds + ) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): from litellm.types.utils import StandardLoggingPayload @@ -1815,6 +1861,50 @@ class PrometheusLogger(CustomLogger): ) return + def _record_guardrail_metrics( + self, + guardrail_name: str, + latency_seconds: float, + status: str, + error_type: Optional[str], + hook_type: str, + ): + """ + Record guardrail metrics for prometheus. + + Args: + guardrail_name: Name of the guardrail + latency_seconds: Execution latency in seconds + status: "success" or "error" + error_type: Type of error if any, None otherwise + hook_type: "pre_call", "during_call", or "post_call" + """ + try: + # Record latency + self.litellm_guardrail_latency_metric.labels( + guardrail_name=guardrail_name, + status=status, + error_type=error_type or "none", + hook_type=hook_type, + ).observe(latency_seconds) + + # Record request count + self.litellm_guardrail_requests_total.labels( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ).inc() + + # Record error count if there was an error + if status == "error" and error_type: + self.litellm_guardrail_errors_total.labels( + guardrail_name=guardrail_name, + error_type=error_type, + hook_type=hook_type, + ).inc() + except Exception as e: + verbose_logger.debug(f"Error recording guardrail metrics: {str(e)}") + @staticmethod def _get_exception_class_name(exception: Exception) -> str: exception_class_name = "" diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c331e4303..d8e8219927 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1645,9 +1645,12 @@ def convert_to_anthropic_tool_result( ) elif content["type"] == "image_url": format = content["image_url"].get("format") if isinstance(content["image_url"], dict) else None - anthropic_content_list.append( - create_anthropic_image_param(content["image_url"], format=format) + _anthropic_image_param = create_anthropic_image_param(content["image_url"], format=format) + _anthropic_image_param = add_cache_control_to_content( + anthropic_content_element=_anthropic_image_param, + original_content_element=content, ) + anthropic_content_list.append(_anthropic_image_param) anthropic_content = anthropic_content_list anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f2b2d7da16..61bffde3ac 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -294,7 +294,11 @@ class ProxyBaseLLMRequestProcessing: if response_cost is not None: try: # Convert response_cost to float if it's a string - cost_value = float(response_cost) if isinstance(response_cost, str) else response_cost + cost_value = ( + float(response_cost) + if isinstance(response_cost, str) + else response_cost + ) if cost_value > 0: updated_spend = current_spend + cost_value except (ValueError, TypeError): @@ -433,6 +437,16 @@ class ProxyBaseLLMRequestProcessing: ) -> Tuple[dict, LiteLLMLoggingObj]: start_time = datetime.now() # start before calling guardrail hooks + # Calculate request queue time if arrival_time is available + # Use start_time.timestamp() to avoid extra time.time() call for better performance + proxy_server_request = self.data.get("proxy_server_request", {}) + arrival_time = proxy_server_request.get("arrival_time") + queue_time_seconds = None + if arrival_time is not None: + # Convert start_time (datetime) to timestamp for calculation + processing_start_time = start_time.timestamp() + queue_time_seconds = processing_start_time - arrival_time + self.data = await add_litellm_data_to_request( data=self.data, request=request, @@ -442,6 +456,19 @@ class ProxyBaseLLMRequestProcessing: proxy_config=proxy_config, ) + # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved + if queue_time_seconds is not None: + from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name + + _metadata_variable_name = _get_metadata_variable_name(request) + if _metadata_variable_name not in self.data: + self.data[_metadata_variable_name] = {} + if not isinstance(self.data[_metadata_variable_name], dict): + self.data[_metadata_variable_name] = {} + self.data[_metadata_variable_name][ + "queue_time_seconds" + ] = queue_time_seconds + self.data["model"] = ( general_settings.get("completion_model", None) # server default or user_model # model name passed via cli args @@ -1235,9 +1262,9 @@ class ProxyBaseLLMRequestProcessing: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs["cache_creation_input_tokens"] = ( - cache_creation_input_tokens - ) + usage_kwargs[ + "cache_creation_input_tokens" + ] = cache_creation_input_tokens if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 5b5723efc3..3f844f21eb 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -161,7 +161,6 @@ class KeyAndTeamLoggingSettings: @staticmethod def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): - if ( user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata @@ -174,12 +173,12 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - ) - team_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) - ) + key_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + team_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) ######################################################################################### # Key-based callbacks ######################################################################################### @@ -462,7 +461,6 @@ class LiteLLMProxyRequestSetup: team_id=user_api_key_dict.team_id, ) # handles aliases, wildcards, etc. ): - _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call( headers, user_api_key_dict ) @@ -663,11 +661,11 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = ( - LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], - ) + data[_metadata_variable_name][ + "tags" + ] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], ) if "disable_global_guardrails" in key_metadata and isinstance( key_metadata["disable_global_guardrails"], bool @@ -815,11 +813,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # Init - Proxy Server Request # we do this as soon as entering so we track the original request ########################################################## + # Track arrival time for queue time metric + arrival_time = time.time() data["proxy_server_request"] = { "url": str(request.url), "method": request.method, "headers": _headers, "body": copy.copy(data), # use copy instead of deepcopy + "arrival_time": arrival_time, # Track when request arrived at proxy } safe_add_api_version_from_query_params(data, request) @@ -930,9 +931,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name]["global_max_parallel_requests"] = ( - general_settings.get("global_max_parallel_requests", None) - ) + data[_metadata_variable_name][ + "global_max_parallel_requests" + ] = general_settings.get("global_max_parallel_requests", None) ### KEY-LEVEL Controls key_metadata = user_api_key_dict.metadata diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0929d08875..171898b163 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -962,6 +962,7 @@ class ProxyLogging: Updated data dictionary if guardrail passes, None if guardrail should be skipped """ from litellm.types.guardrails import GuardrailEventHooks + from litellm.integrations.prometheus import PrometheusLogger # Determine the event type based on call type event_type = GuardrailEventHooks.pre_call @@ -974,30 +975,62 @@ class ProxyLogging: guardrail_name = callback.guardrail_name - # Check if load balancing should be used - if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): - response = await self._execute_guardrail_with_load_balancing( - guardrail_name=guardrail_name, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) - else: - # Single guardrail - execute directly - response = await self._execute_guardrail_hook( - callback=callback, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) + # Track timing and errors for prometheus metrics + # Use time.perf_counter() for more accurate duration measurements + guardrail_start_time = time.perf_counter() + status = "success" + error_type = None - # Process the response if one was returned - if response is not None: - data = await self.process_pre_call_hook_response( - response=response, data=data, call_type=call_type - ) + try: + # Check if load balancing should be used + if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): + response = await self._execute_guardrail_with_load_balancing( + guardrail_name=guardrail_name, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + else: + # Single guardrail - execute directly + response = await self._execute_guardrail_hook( + callback=callback, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + + # Process the response if one was returned + if response is not None: + data = await self.process_pre_call_hook_response( + response=response, data=data, call_type=call_type + ) + + except Exception as e: + status = "error" + error_type = type(e).__name__ + # Re-raise the exception to maintain existing behavior + raise + finally: + # Record prometheus metrics + guardrail_end_time = time.perf_counter() + latency_seconds = guardrail_end_time - guardrail_start_time + + # Get guardrail name for metrics (fallback if not set) + metrics_guardrail_name = guardrail_name or getattr(callback, "guardrail_name", callback.__class__.__name__) or "unknown" + + # Find PrometheusLogger in callbacks and record metrics + for prom_callback in litellm.callbacks: + if isinstance(prom_callback, PrometheusLogger): + prom_callback._record_guardrail_metrics( + guardrail_name=metrics_guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type="pre_call", + ) + break return data diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index a92b5d25a3..7667d1bad8 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -443,11 +443,18 @@ class ResponseAPILoggingUtils: completion_tokens=0, total_tokens=0, ) - response_api_usage: ResponseAPIUsage = ( - ResponseAPIUsage(**usage_input) - if isinstance(usage_input, dict) - else usage_input - ) + response_api_usage: ResponseAPIUsage + if isinstance(usage_input, dict): + total_tokens = usage_input.get("total_tokens") + if total_tokens is None: + input_tokens = usage_input.get("input_tokens") + output_tokens = usage_input.get("output_tokens") + if input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + usage_input["total_tokens"] = total_tokens + response_api_usage = ResponseAPIUsage(**usage_input) + else: + response_api_usage = usage_input prompt_tokens: int = response_api_usage.input_tokens or 0 completion_tokens: int = response_api_usage.output_tokens or 0 prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index fb439a9541..88dee19ae5 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -185,6 +185,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_redis_daily_spend_update_queue_size", "litellm_in_memory_spend_update_queue_size", "litellm_redis_spend_update_queue_size", + "litellm_request_queue_time_seconds", + "litellm_guardrail_latency_seconds", + "litellm_guardrail_errors_total", + "litellm_guardrail_requests_total", # Cache metrics "litellm_cache_hits_metric", "litellm_cache_misses_metric", @@ -223,6 +227,23 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, ] + litellm_request_queue_time_seconds = [ + UserAPIKeyLabelNames.END_USER.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.USER.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + # Guardrail metrics - these use custom labels (guardrail_name, status, error_type, hook_type) + # which are not part of UserAPIKeyLabelNames + litellm_guardrail_latency_seconds: List[str] = [] + litellm_guardrail_errors_total: List[str] = [] + litellm_guardrail_requests_total: List[str] = [] + litellm_proxy_total_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, @@ -479,11 +500,6 @@ class PrometheusMetricLabels: return default_labels + custom_labels -from typing import List, Optional - -from pydantic import BaseModel, Field - - class UserAPIKeyLabelValues(BaseModel): end_user: Annotated[ Optional[str], Field(..., alias=UserAPIKeyLabelNames.END_USER.value) diff --git a/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py b/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py new file mode 100644 index 0000000000..0743a9c7ba --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py @@ -0,0 +1,424 @@ +""" +Unit tests for prometheus queue time and guardrail metrics +""" +from datetime import datetime +from unittest.mock import MagicMock + +import pytest +from prometheus_client import REGISTRY + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + """Clean up prometheus registry between tests""" + # Clear the registry before each test + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + yield + # Clean up after test + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + + +class TestPrometheusQueueTimeMetric: + """Test request queue time metric recording""" + + def test_queue_time_metric_recorded_in_set_latency_metrics(self): + """Test that queue time metric is recorded when queue_time_seconds is present in metadata""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs with queue_time_seconds in metadata + queue_time_seconds = 0.5 + + kwargs = { + "litellm_params": {"metadata": {"queue_time_seconds": queue_time_seconds}}, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should be called + mock_metric.labels.assert_called() + # Check that observe was called on the queue time metric + assert mock_labeled_metric.observe.called + # Verify the observed value + observed_value = None + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0: + observed_value = call[0][0] + if observed_value == queue_time_seconds: + break + assert observed_value == queue_time_seconds + assert observed_value >= 0 + + def test_queue_time_metric_not_recorded_when_missing(self): + """Test that queue time metric is not recorded when queue_time_seconds is missing""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs without queue_time_seconds + kwargs = { + "litellm_params": {"metadata": {}}, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should not be called (queue_time_seconds is None) + # We check that observe was not called with queue_time_seconds + queue_time_called = False + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0 and call[0][0] == 0.5: # Our test queue time value + queue_time_called = True + break + assert ( + not queue_time_called + ), "Queue time metric should not be recorded when queue_time_seconds is missing" + + def test_queue_time_metric_not_recorded_when_negative(self): + """Test that queue time metric is not recorded when queue_time_seconds is negative""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs with negative queue_time_seconds + kwargs = { + "litellm_params": { + "metadata": {"queue_time_seconds": -0.1} # Negative value + }, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should not be called for negative values + # We check that observe was not called with the negative value + negative_value_called = False + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0 and call[0][0] == -0.1: + negative_value_called = True + break + assert ( + not negative_value_called + ), "Queue time metric should not be recorded for negative values" + + +class TestPrometheusGuardrailMetrics: + """Test guardrail metrics recording""" + + def test_record_guardrail_metrics_success(self): + """Test recording guardrail metrics for successful execution""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + mock_errors_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric + + guardrail_name = "test_guardrail" + latency_seconds = 0.15 + status = "success" + error_type = None + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + + # Assert - latency metric should be recorded + mock_latency_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + error_type="none", + hook_type=hook_type, + ) + mock_latency_metric.labels.return_value.observe.assert_called_once_with( + latency_seconds + ) + + # Assert - requests metric should be incremented + mock_requests_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ) + mock_requests_metric.labels.return_value.inc.assert_called_once() + + # Assert - errors metric should NOT be called for success + mock_errors_metric.labels.assert_not_called() + + def test_record_guardrail_metrics_error(self): + """Test recording guardrail metrics for failed execution""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + mock_errors_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric + + guardrail_name = "test_guardrail" + latency_seconds = 0.2 + status = "error" + error_type = "ValueError" + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + + # Assert - latency metric should be recorded + mock_latency_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + mock_latency_metric.labels.return_value.observe.assert_called_once_with( + latency_seconds + ) + + # Assert - requests metric should be incremented + mock_requests_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ) + mock_requests_metric.labels.return_value.inc.assert_called_once() + + # Assert - errors metric should be incremented + mock_errors_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + error_type=error_type, + hook_type=hook_type, + ) + mock_errors_metric.labels.return_value.inc.assert_called_once() + + def test_record_guardrail_metrics_during_call_hook(self): + """Test recording guardrail metrics for during_call hook""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + + guardrail_name = "moderation_guardrail" + latency_seconds = 0.1 + status = "success" + hook_type = "during_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=None, + hook_type=hook_type, + ) + + # Assert - hook_type should be "during_call" + mock_latency_metric.labels.assert_called_once() + call_kwargs = mock_latency_metric.labels.call_args[1] + assert call_kwargs["hook_type"] == "during_call" + + def test_record_guardrail_metrics_handles_exception(self): + """Test that _record_guardrail_metrics handles exceptions gracefully""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metric to raise exception + mock_metric = MagicMock() + mock_metric.labels.side_effect = Exception("Test error") + prometheus_logger.litellm_guardrail_latency_metric = mock_metric + prometheus_logger.litellm_guardrail_requests_total = MagicMock() + + # Act & Assert - should not raise exception + try: + prometheus_logger._record_guardrail_metrics( + guardrail_name="test", + latency_seconds=0.1, + status="success", + error_type=None, + hook_type="pre_call", + ) + except Exception: + pytest.fail("_record_guardrail_metrics should handle exceptions gracefully") + + def test_record_guardrail_metrics_with_guardrail_name_attribute(self): + """Test that guardrail name is extracted from guardrail_name attribute if available""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + + guardrail_name = "custom_guardrail_name" + latency_seconds = 0.1 + status = "success" + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=None, + hook_type=hook_type, + ) + + # Assert - guardrail_name should be used + mock_latency_metric.labels.assert_called_once() + call_kwargs = mock_latency_metric.labels.call_args[1] + assert call_kwargs["guardrail_name"] == guardrail_name diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 4914ec0bfb..42a2b5d097 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -620,26 +620,26 @@ def test_bedrock_tools_unpack_defs(): def test_bedrock_image_processor_content_type_fallback_url_extension(): """ - Test that _post_call_image_processing falls back to URL extension + Test that _post_call_image_processing falls back to URL extension when content-type is binary/octet-stream or application/octet-stream """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a simple PNG header (magic bytes) png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" png_content = png_header + b"\x00" * 100 # Add some padding mock_response.content = png_content - + # Test with .png URL image_url = "https://example.com/test-image.png" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/png" assert base64_bytes == base64.b64encode(png_content).decode("utf-8") @@ -650,22 +650,22 @@ def test_bedrock_image_processor_content_type_fallback_binary_detection(): when content-type is missing and URL extension is not recognized """ import base64 - + # Create mock response with no content-type mock_response = MagicMock() mock_response.headers.get.return_value = None - + # Create a JPEG header (magic bytes) jpeg_header = b"\xff\xd8\xff" jpeg_content = jpeg_header + b"\x00" * 100 # Add some padding mock_response.content = jpeg_content - + # Test with URL without extension image_url = "https://example.com/test-image-without-extension" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/jpeg" assert base64_bytes == base64.b64encode(jpeg_content).decode("utf-8") @@ -675,22 +675,22 @@ def test_bedrock_image_processor_content_type_fallback_application_octet_stream( Test that _post_call_image_processing handles application/octet-stream correctly """ import base64 - + # Create mock response with application/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "application/octet-stream" - + # Create a GIF header (magic bytes) gif_header = b"GIF8" + b"\x00" + b"a" gif_content = gif_header + b"\x00" * 100 # Add some padding mock_response.content = gif_content - + # Test with .gif URL image_url = "https://s3.amazonaws.com/bucket/image.gif" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/gif" assert base64_bytes == base64.b64encode(gif_content).decode("utf-8") @@ -700,22 +700,22 @@ def test_bedrock_image_processor_content_type_with_query_params(): Test that _post_call_image_processing correctly extracts extension from URL with query parameters """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a WebP header (magic bytes) webp_header = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" webp_content = webp_header + b"\x00" * 100 # Add some padding mock_response.content = webp_content - + # Test with URL containing query parameters (common in S3 signed URLs) image_url = "https://s3.amazonaws.com/bucket/image.webp?AWSAccessKeyId=123&Expires=456&Signature=789" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/webp" assert base64_bytes == base64.b64encode(webp_content).decode("utf-8") @@ -725,21 +725,21 @@ def test_bedrock_image_processor_content_type_normal_header(): Test that _post_call_image_processing works normally when content-type is correctly set """ import base64 - + # Create mock response with correct content-type mock_response = MagicMock() mock_response.headers.get.return_value = "image/png" - + # Create a PNG header png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" png_content = png_header + b"\x00" * 100 mock_response.content = png_content - + image_url = "https://example.com/test-image.png" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/png" assert base64_bytes == base64.b64encode(png_content).decode("utf-8") @@ -751,16 +751,16 @@ def test_bedrock_image_processor_content_type_fallback_failure(): # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create content with unrecognizable image format mock_response.content = b"\x00" * 100 - + # Test with URL without recognizable extension image_url = "https://example.com/unknown-file" - + with pytest.raises(ValueError) as excinfo: BedrockImageProcessor._post_call_image_processing(mock_response, image_url) - + assert "Unable to determine content type" in str(excinfo.value) @@ -771,18 +771,18 @@ def test_bedrock_image_processor_content_type_jpeg_variants(): # Create mock response with binary/octet-stream mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + jpeg_header = b"\xff\xd8\xff" jpeg_content = jpeg_header + b"\x00" * 100 mock_response.content = jpeg_content - + # Test with .jpg extension image_url_jpg = "https://example.com/photo.jpg" _, content_type_jpg = BedrockImageProcessor._post_call_image_processing( mock_response, image_url_jpg ) assert content_type_jpg == "image/jpeg" - + # Test with .jpeg extension image_url_jpeg = "https://example.com/photo.jpeg" _, content_type_jpeg = BedrockImageProcessor._post_call_image_processing( @@ -797,22 +797,22 @@ def test_bedrock_image_processor_content_type_pdf_document(): when content-type is binary/octet-stream """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a PDF header (magic bytes: %PDF) pdf_header = b"%PDF-1.4" pdf_content = pdf_header + b"\x00" * 100 mock_response.content = pdf_content - + # Test with .pdf URL pdf_url = "https://s3.amazonaws.com/bucket/document.pdf" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, pdf_url ) - + assert content_type == "application/pdf" assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") @@ -822,12 +822,12 @@ def test_bedrock_image_processor_content_type_document_formats(): Test that _post_call_image_processing handles various document formats """ import base64 - + # Create mock response mock_response = MagicMock() mock_response.headers.get.return_value = "application/octet-stream" mock_response.content = b"\x00" * 100 - + # Test various document formats test_cases = [ ("https://example.com/doc.pdf", "application/pdf"), @@ -837,7 +837,7 @@ def test_bedrock_image_processor_content_type_document_formats(): ("https://example.com/page.html", "text/html"), ("https://example.com/readme.txt", "text/plain"), ] - + for url, expected_mime in test_cases: _, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, url @@ -850,21 +850,21 @@ def test_bedrock_image_processor_content_type_s3_pdf_with_query(): Test that _post_call_image_processing handles S3 PDF with query parameters """ import base64 - + # Create mock response mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + pdf_content = b"%PDF-1.4" + b"\x00" * 100 mock_response.content = pdf_content - + # S3 signed URL with query parameters s3_url = "https://my-bucket.s3.us-east-1.amazonaws.com/documents/report.pdf?AWSAccessKeyId=AKIAIOSFODNN7EXAMPLE&Expires=1234567890&Signature=abcdef123456" - + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, s3_url ) - + assert content_type == "application/pdf" assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") @@ -1139,6 +1139,170 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["format"] == format_type +def test_convert_to_anthropic_tool_result_image_with_cache_control(): + """ + Test that cache_control is properly applied to image content in tool results. + This tests the functionality added in the uncommitted changes where + add_cache_control_to_content is called for image_url content types. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + # Test with base64 image data URI + message = { + "role": "tool", + "tool_call_id": "call_test_123", + "content": [ + { + "type": "text", + "text": "Here is the image you requested:", + }, + { + "type": "image_url", + "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQ", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify the result structure + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "call_test_123" + assert isinstance(result["content"], list) + assert len(result["content"]) == 2 + + # Verify text content + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Here is the image you requested:" + + # Verify image content with cache_control + assert result["content"][1]["type"] == "image" + assert result["content"][1]["source"]["type"] == "base64" + assert result["content"][1]["source"]["media_type"] == "image/jpeg" + assert "cache_control" in result["content"][1] + assert result["content"][1]["cache_control"]["type"] == "ephemeral" + + +def test_convert_to_anthropic_tool_result_image_without_cache_control(): + """ + Test that images without cache_control in tool results work correctly. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_test_456", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA", + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify the result structure + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "call_test_456" + assert isinstance(result["content"], list) + assert len(result["content"]) == 1 + + # Verify image content without cache_control (cache_control will be None if not set) + assert result["content"][0]["type"] == "image" + assert result["content"][0]["source"]["type"] == "base64" + assert result["content"][0]["source"]["media_type"] == "image/png" + assert result["content"][0].get("cache_control") is None + + +def test_convert_to_anthropic_tool_result_mixed_content_with_cache_control(): + """ + Test tool results with mixed content types (text and image) where only some have cache_control. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_test_789", + "content": [ + { + "type": "text", + "text": "First image:", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "image_url", + "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRg", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "text", + "text": "Second image (no cache):", + }, + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + assert result["type"] == "tool_result" + assert isinstance(result["content"], list) + assert len(result["content"]) == 4 + + # First text with cache_control + assert result["content"][0]["type"] == "text" + assert result["content"][0]["cache_control"]["type"] == "ephemeral" + + # First image with cache_control + assert result["content"][1]["type"] == "image" + assert result["content"][1]["cache_control"]["type"] == "ephemeral" + + # Second text without cache_control (cache_control will be None if not set) + assert result["content"][2]["type"] == "text" + assert result["content"][2].get("cache_control") is None + + # Second image without cache_control (cache_control will be None if not set) + assert result["content"][3]["type"] == "image" + assert result["content"][3].get("cache_control") is None + + +def test_convert_to_anthropic_tool_result_image_url_as_http(): + """ + Test that HTTP/HTTPS URLs with cache_control are handled correctly. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_http_001", + "content": [ + { + "type": "image_url", + "image_url": "https://example.com/image.jpg", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify image is passed as URL reference with cache_control + assert result["content"][0]["type"] == "image" + assert result["content"][0]["source"]["type"] == "url" + assert result["content"][0]["source"]["url"] == "https://example.com/image.jpg" + assert result["content"][0]["cache_control"]["type"] == "ephemeral" def test_anthropic_messages_pt_server_tool_use_passthrough(): """ Test that anthropic_messages_pt passes through server_tool_use and diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 96ac2e2c34..09628cd4a7 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -203,3 +203,22 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens == 0 assert result.completion_tokens == 20 assert result.total_tokens == 20 + + def test_transform_response_api_usage_calculates_total_from_input_and_output_tokens_if_available(self): + """Test transformation calculates total_tokens when it's None and input / output tokens are present""" + # Setup + usage = { + "input_tokens": 15, + "output_tokens": 25, + "total_tokens": None, + } + + # Execute + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + + # Assert + assert result.prompt_tokens == 15 + assert result.completion_tokens == 25 + assert result.total_tokens == 40 # 15 + 25