tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-11 17:05:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18715 from BerriAI/litellm_staging_01_06_2026

Browse Source 
Staging 01/06/2026
This commit is contained in:
Sameer Kankute
2026-01-09 17:10:04 +05:30
committed by GitHub
parent bb9347207b e7efd51bd7
commit 072e3dcd0c
10 changed files with 879 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/prometheus.py
+90
View File
@@ -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 = ""
litellm/litellm_core_utils/prompt_templates/factory.py
+5 -2
View File
@@ -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
litellm/proxy/common_request_processing.py
+31 -4
View File
@@ -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
litellm/proxy/litellm_pre_call_utils.py
+17 -16
View File
@@ -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
litellm/proxy/utils.py
+56 -23
View File
@@ -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
litellm/responses/utils.py
+12 -5
View File
@@ -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
litellm/types/integrations/prometheus.py
+21 -5
View File
@@ -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)
tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py
+424
View File
@@ -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
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
+204 -40
View File
@@ -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
tests/test_litellm/responses/test_responses_utils.py
+19
View File
@@ -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