scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_fix_image_scan_presets_import

Browse source
This commit is contained in:
mateo-berri 2026-09-02 20:57:07 -07:00
commit 06d1bf1b3f
95 changed files with 6403 additions and 377 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
.github/workflows/test-code-quality.yml vendored
View file

@ -128,6 +128,9 @@ jobs:
- name: check_fastuuid_usage
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
- name: check_py310_typing_imports
run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py
- name: check_e2e_no_raw_requests
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
@ -145,3 +148,33 @@ jobs:
- name: documentation_test_api_docs
run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
python-310-import-smoke:
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.10"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: uv sync --frozen --extra proxy --python 3.10
- run: uv run --no-sync python --version
- name: Import litellm
run: uv run --no-sync python -c "import litellm"
- name: Check litellm CLI
run: uv run --no-sync litellm --version

6
basedpyright-code-budget.json
View file

@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15290
"limit": 15288
},
"reportMissingTypeStubs": {
"limit": 40
@ -105,7 +105,7 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38332
"limit": 38324
},
"reportUnknownParameterType": {
"limit": 19625
@ -123,7 +123,7 @@
"limit": 4
},
"reportUnnecessaryIsInstance": {
"limit": 823
"limit": 819
},
"reportUntypedBaseClass": {
"limit": 0

188
litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs
View file

@ -186,68 +186,180 @@ mod tests {
use serde_json::json;
#[test]
fn supported_params_match_python_mistral_ocr_config() {
fn extract_header_is_a_supported_ocr_param() {
assert!(supported_ocr_params().contains(&"extract_header"));
}
#[test]
fn extract_footer_is_a_supported_ocr_param() {
assert!(supported_ocr_params().contains(&"extract_footer"));
}
#[test]
fn existing_ocr_params_remain_supported() {
for param in [
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
] {
assert!(supported_ocr_params().contains(&param));
}
}
#[test]
fn map_ocr_params_forwards_extract_header() {
let params = json!({"extract_header": true});
assert_eq!(
supported_ocr_params(),
&[
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
"document_annotation_prompt",
"extract_header",
"extract_footer",
"table_format",
"confidence_scores_granularity",
"include_blocks",
"id",
]
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_forwards_extract_footer() {
let params = json!({"extract_footer": true});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_forwards_extract_header_and_footer() {
let params = json!({"extract_header": true, "extract_footer": false});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_drops_unknown_params() {
let params = json!({
"extract_header": true,
"unsupported_param": "value",
"pages": [0, 1]
});
let params = json!({"extract_header": true, "unsupported_param": "value"});
let mapped = map_ocr_params(params.as_object().unwrap());
assert_eq!(mapped.get("extract_header"), Some(&json!(true)));
assert_eq!(mapped.get("pages"), Some(&json!([0, 1])));
assert!(!mapped.contains_key("unsupported_param"));
}
#[test]
fn transform_ocr_request_builds_mistral_body() {
fn new_ocr_params_are_supported() {
for param in [
"table_format",
"confidence_scores_granularity",
"document_annotation_prompt",
"include_blocks",
"id",
] {
assert!(supported_ocr_params().contains(&param));
}
}
#[test]
fn map_ocr_params_forwards_new_ocr_params() {
for (param, value) in [
("table_format", json!("html")),
("confidence_scores_granularity", json!("word")),
(
"document_annotation_prompt",
json!("Extract all invoice line items"),
),
("include_blocks", json!(true)),
("id", json!("req-123")),
] {
let params = json!({param: value});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
}
#[test]
fn transform_ocr_request_includes_each_optional_param() {
let document = json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
});
for (param, value) in [
("table_format", json!("html")),
("confidence_scores_granularity", json!("word")),
(
"document_annotation_prompt",
json!("Extract all invoice line items"),
),
("id", json!("req-123")),
("extract_header", json!(true)),
("include_blocks", json!(true)),
("pages", json!([0, 1])),
] {
let result = transform_ocr_request(
"mistral-ocr-latest",
document.clone(),
json!({param: value}).as_object().unwrap().clone(),
)
.expect("request should transform");
assert_eq!(result.data.get(param), Some(&value));
assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest")));
assert_eq!(result.data.get("document"), Some(&document));
assert_eq!(result.files, None);
}
}
#[test]
fn transform_ocr_request_includes_multiple_new_params() {
let document = json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
});
let optional_params = json!({
"include_image_base64": true,
"table_format": "html"
"table_format": "html",
"confidence_scores_granularity": "page",
"extract_header": true
})
.as_object()
.unwrap()
.clone();
let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params)
let result = transform_ocr_request("mistral-ocr-latest", document, optional_params)
.expect("request should transform");
assert_eq!(result.data.get("table_format"), Some(&json!("html")));
assert_eq!(
result.data,
json!({
"model": "mistral-ocr-latest",
"document": document,
"include_image_base64": true,
"table_format": "html"
})
result.data.get("confidence_scores_granularity"),
Some(&json!("page"))
);
assert_eq!(result.files, None);
assert_eq!(result.data.get("extract_header"), Some(&json!(true)));
}
#[test]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let blocks = json!([{"type": "title", "content": "Invoice"}]);
let confidence_scores = json!({"page": 0.98});
let response = json!({
"pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}],
"model": "mistral-ocr-4-0",
"usage_info": {"pages_processed": 1}
});
let result =
transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform");
assert_eq!(result.pages[0].get("blocks"), Some(&blocks));
assert_eq!(
result.pages[0].get("confidence_scores"),
Some(&confidence_scores)
);
}
#[test]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let response = json!({
"pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}],
"model": "mistral-ocr-4-0",
"usage_info": {"pages_processed": 1}
});
let result = transform_ocr_response("mistral-ocr-4-0", response.clone())
.expect("response should transform");
assert_eq!(result.pages[0], response["pages"][0]);
}
#[test]

1
litellm/constants.py
View file

@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
"Truncation is a DB storage safeguard. "

49
litellm/integrations/azure_storage/azure_storage.py
View file

@ -1,7 +1,9 @@
import asyncio
import os
import time
from collections.abc import Callable
from datetime import datetime, timedelta
from functools import cache
from typing import Final
from litellm._logging import verbose_logger
@ -19,21 +21,40 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.secret_managers.get_azure_ad_token_provider import (
get_azure_ad_token_provider,
)
from litellm.types.secret_managers.get_azure_ad_token_provider import (
AzureCredentialType,
)
from litellm.types.utils import StandardLoggingPayload
AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
@cache
def _cached_credential_chain_token_provider() -> Callable[[], str]:
return get_azure_ad_token_provider(
azure_scope=AZURE_STORAGE_TOKEN_SCOPE,
azure_credential=AzureCredentialType.DefaultAzureCredential,
)
class AzureBlobStorageLogger(CustomBatchLogger):
def __init__(
self,
build_credential_chain_token_provider: Callable[
[], Callable[[], str]
] = _cached_credential_chain_token_provider,
**kwargs,
):
try:
verbose_logger.debug("AzureBlobStorageLogger: in init azure blob storage logger")
# Env Variables used for Azure Storage Authentication
self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID")
self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID")
self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET")
self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID") or None
self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID") or None
self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET") or None
self.azure_storage_account_key: str | None = os.getenv("AZURE_STORAGE_ACCOUNT_KEY")
# Required Env Variables for Azure Storage
@ -55,6 +76,9 @@ class AzureBlobStorageLogger(CustomBatchLogger):
# Internal variables used for Token based authentication
self.azure_auth_token: str | None = None # the Azure AD token to use for Azure Storage API requests
self.token_expiry: datetime | None = None # the expiry time of the currentAzure AD token
self._build_credential_chain_token_provider: Callable[[], Callable[[], str]] = (
build_credential_chain_token_provider
)
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
@ -231,10 +255,15 @@ class AzureBlobStorageLogger(CustomBatchLogger):
"""
Wrapper to set self.azure_auth_token to a valid Azure AD token, refreshing if necessary
Refreshes the token when:
- Token is expired
- Token is not set
Without a service principal configured, the credential chain provider is read every
time; it caches internally and refreshes against the token's real expiry. The read runs
in a worker thread because the chain walk (IMDS probe, CLI subprocess) is blocking
"""
if self.tenant_id is None and self.client_id is None and self.client_secret is None:
token_provider: Final = self._build_credential_chain_token_provider()
self.azure_auth_token = await asyncio.to_thread(token_provider)
return
# Check if token needs refresh
if self._azure_ad_token_is_expired() or self.azure_auth_token is None:
verbose_logger.debug("Azure AD token needs refresh")
@ -273,13 +302,9 @@ class AzureBlobStorageLogger(CustomBatchLogger):
tenant_id=tenant_id,
client_id=client_id,
client_secret=client_secret,
scope="https://storage.azure.com/.default",
scope=AZURE_STORAGE_TOKEN_SCOPE,
)
token: Final = token_provider()
verbose_logger.debug("azure auth token %s", token)
return token
return token_provider()
def _azure_ad_token_is_expired(self):
"""

217
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -19,11 +19,13 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.constants import REDACTED_BY_LITELLM
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_handler import (
get_datadog_base_url_from_env,
get_datadog_service,
get_datadog_tags,
normalize_datadog_tag_value,
)
from litellm.integrations.datadog.datadog_mock_client import (
create_mock_datadog_client,
@ -34,6 +36,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
handle_any_messages_to_chat_completion_str_messages_conversion,
)
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.custom_httpx.http_handler import (
@ -43,6 +46,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
from litellm.types.integrations.datadog_llm_obs import *
from litellm.types.utils import (
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
CallTypes,
StandardLoggingGuardrailInformation,
StandardLoggingPayload,
@ -52,6 +56,120 @@ from litellm.types.utils import (
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
_SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset(
{"agent", "assistant", "developer", "function", "model", "system", "tool", "user"}
)
_PROMPT_CARRYING_METADATA_FIELDS: Final = frozenset(
{
"routing_decision",
"requester_metadata",
"prompt_management_metadata",
"mcp_tool_call_metadata",
"vector_store_request_metadata",
}
)
_ROUTER_SPAN_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{
"tier": "router_tier",
"cause": "router_cause",
"score": "router_score",
"escalated": "router_escalated",
"signals": "router_signals",
"routed_model": "routed_model",
}
)
_ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "router_escalated", "routed_model")
_COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS)
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]:
metadata: Final = standard_logging_payload.get("metadata")
return metadata or _EMPTY_MAPPING
def _router_span_fields(
standard_logging_payload: StandardLoggingPayload, redact_prompt_text: bool
) -> Mapping[str, object]:
"""Flatten the auto-router decision, omitting prompt-quoting fields when redaction is enabled."""
routing_decision: Final = _mapping_field(_metadata_of(standard_logging_payload), "routing_decision")
if not routing_decision:
return _EMPTY_MAPPING
escalated: Final = bool(routing_decision.get("escalated") or routing_decision.get("context_escalated"))
return MappingProxyType(
{
_ROUTER_SPAN_FIELDS[record_field]: value
for record_field, value in (*routing_decision.items(), ("escalated", escalated))
if record_field in _ROUTER_SPAN_FIELDS
and value is not None
and not (redact_prompt_text and record_field in PROMPT_QUOTING_ROUTING_DECISION_FIELDS)
}
)
def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]:
"""The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text."""
return MappingProxyType(
{
field: value
for field, value in standard_logging_metadata.items()
if field not in _PROMPT_CARRYING_METADATA_FIELDS
}
)
def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]:
"""Each message's shape with its content replaced and tool payloads dropped; no message is invented."""
return tuple(
{
"role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "",
"content": REDACTED_BY_LITELLM,
}
for message in messages
for role in (message.get("role", ""),)
)
def _cost_dimension_tags(
standard_logging_payload: StandardLoggingPayload, router_fields: Mapping[str, object]
) -> tuple[str, ...]:
"""The dimensions LLM Obs breaks token and cost metrics down by, as span tags."""
metadata: Final = _metadata_of(standard_logging_payload)
dimensions: Final = (
("user", metadata.get("user_api_key_user_id")),
("key_alias", metadata.get("user_api_key_alias")),
("model_group", standard_logging_payload.get("model_group")),
*((dimension, router_fields.get(dimension)) for dimension in _ROUTER_DIMENSIONS),
)
return tuple(
f"{key}:{normalized}"
for key, value in dimensions
if value is not None and (normalized := normalize_datadog_tag_value(value)) != ""
)
def _declared_cost_tags(span_tags: Sequence[str]) -> tuple[str, ...]:
"""Declare only cost dimensions carrying a value on this span."""
present: Final = frozenset(key for tag in span_tags if (key := tag.partition(":")[0]) and tag.partition(":")[2])
return tuple(dimension for dimension in _COST_DIMENSIONS if dimension in present)
def _reasoning_output_tokens(usage_object: Mapping[str, Any] | None) -> float:
"""The provider's reasoning-token count, from either the chat or the responses spelling."""
if usage_object is None:
return 0.0
return next(
(
float(reasoning_tokens)
for details_field in ("completion_tokens_details", "output_tokens_details")
if isinstance(
reasoning_tokens := _mapping_field(usage_object, details_field).get("reasoning_tokens"), (int, float)
)
and not isinstance(reasoning_tokens, bool)
),
0.0,
)
def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]:
@ -316,12 +434,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
dict_datadog_llm_obs_params: dict = {}
if litellm.datadog_llm_observability_params is not None:
if isinstance(litellm.datadog_llm_observability_params, DatadogLLMObsInitParams):
dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump()
dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump(exclude_unset=True)
elif isinstance(litellm.datadog_llm_observability_params, dict):
# only allow params that are of DatadogLLMObsInitParams
dict_datadog_llm_obs_params = DatadogLLMObsInitParams(
**litellm.datadog_llm_observability_params
).model_dump()
).model_dump(exclude_unset=True)
return dict_datadog_llm_obs_params
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
@ -410,25 +528,40 @@ class DataDogLLMObsLogger(CustomBatchLogger):
if standard_logging_payload is None:
raise Exception("DataDogLLMObs: standard_logging_object is not set")
metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
raw_metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
metadata: Final = raw_metadata if isinstance(raw_metadata, dict) else {}
redact_payload: Final = self._payload_logging_is_off(kwargs)
input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"]))
input_messages: Final = _to_dd_messages(standard_logging_payload.get("messages"))
output_messages: Final = self._get_response_messages(
standard_logging_payload=standard_logging_payload,
call_type=standard_logging_payload.get("call_type"),
)
input_meta: Final = InputMeta(messages=_redact_messages(input_messages) if redact_payload else input_messages)
output_meta: Final = OutputMeta(
messages=self._get_response_messages(
standard_logging_payload=standard_logging_payload,
call_type=standard_logging_payload.get("call_type"),
)
messages=_redact_messages(output_messages) if redact_payload else output_messages
)
error_info: Final = self._assemble_error_info(standard_logging_payload)
metadata_parent_id: str | None = None
if isinstance(metadata, dict):
metadata_parent_id = metadata.get("parent_id")
raw_parent_id: Final = metadata.get("parent_id")
metadata_parent_id: Final[str | None] = str(raw_parent_id) if raw_parent_id else None
tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
tool_definitions: Final = (
() if redact_payload else _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
)
span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id)
payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload)
router_fields: Final = _router_span_fields(standard_logging_payload, redact_prompt_text=redact_payload)
span_tags: Final = [
*get_datadog_tags(standard_logging_object=standard_logging_payload),
*_cost_dimension_tags(standard_logging_payload, router_fields),
]
payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(
standard_logging_payload,
router_fields=router_fields,
cost_tags=_declared_cost_tags(span_tags),
redact_prompt_text=redact_payload,
)
meta: Final[Meta] = {
"kind": span_kind,
@ -451,7 +584,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
duration=int((end_time - start_time).total_seconds() * 1e9),
metrics=metrics,
status="error" if error_info else "ok",
tags=get_datadog_tags(standard_logging_object=standard_logging_payload),
tags=span_tags,
)
apm_trace_id: Final = self._get_apm_trace_id()
@ -497,6 +630,13 @@ class DataDogLLMObsLogger(CustomBatchLogger):
)
return error_info
def _payload_logging_is_off(self, kwargs: Mapping[str, Any]) -> bool:
return (
bool(self.turn_off_message_logging)
or self.message_logging is not True
or should_redact_message_logging(dict(kwargs))
)
def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics:
"""
Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from.
@ -513,10 +653,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
total_cost: Final = float(standard_logging_payload.get("response_cost", 0))
time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload)
raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object")
raw_usage: Final = _metadata_of(standard_logging_payload).get("usage_object")
usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None
cache_read: Final = float(extract_cache_read_tokens(usage_object))
cache_write: Final = float(extract_cache_creation_tokens(usage_object))
reasoning_output_tokens: Final = _reasoning_output_tokens(usage_object)
metrics: Final[LLMMetrics] = {
"input_tokens": prompt_tokens,
@ -533,6 +674,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
if cache_read or cache_write
else {}
),
**({"reasoning_output_tokens": reasoning_output_tokens} if reasoning_output_tokens else {}),
}
return metrics
@ -707,11 +849,21 @@ class DataDogLLMObsLogger(CustomBatchLogger):
# Default fallback for unknown or passthrough operations
return "llm"
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
def _get_dd_llm_obs_payload_metadata(
self,
standard_logging_payload: StandardLoggingPayload,
router_fields: Mapping[str, object] | None = None,
cost_tags: Sequence[str] = (),
redact_prompt_text: bool = False,
) -> dict[str, object]:
"""
Fields to track in DD LLM Observability metadata from litellm standard logging payload
"""
_metadata: Final[dict[str, object]] = {
raw_metadata: Final = _metadata_of(standard_logging_payload)
standard_logging_metadata: Final = (
_metadata_without_prompt_carriers(raw_metadata) if redact_prompt_text else raw_metadata
)
return {
"model_name": standard_logging_payload.get("model", "unknown"),
"model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"),
"id": standard_logging_payload.get("id", "unknown"),
@ -719,26 +871,21 @@ class DataDogLLMObsLogger(CustomBatchLogger):
"cache_hit": standard_logging_payload.get("cache_hit", "unknown"),
"cache_key": standard_logging_payload.get("cache_key", "unknown"),
"saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0),
"guardrail_information": standard_logging_payload.get("guardrail_information", None),
"guardrail_information": (
None if redact_prompt_text else standard_logging_payload.get("guardrail_information", None)
),
"is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload),
"latency_metrics": dict(self._get_latency_metrics(standard_logging_payload)),
"spend_metrics": dict(self._get_spend_metrics(standard_logging_payload)),
**standard_logging_metadata,
**(router_fields or _EMPTY_MAPPING),
**(
{"_dd": {**_mapping_field(standard_logging_metadata, "_dd"), "cost_tags": list(cost_tags)}}
if cost_tags
else _EMPTY_MAPPING
),
}
#########################################################
# Add latency metrics to metadata
#########################################################
latency_metrics: Final = self._get_latency_metrics(standard_logging_payload)
_metadata.update({"latency_metrics": dict(latency_metrics)})
#########################################################
# Add spend metrics to metadata
#########################################################
spend_metrics: Final = self._get_spend_metrics(standard_logging_payload)
_metadata.update({"spend_metrics": dict(spend_metrics)})
_standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {}
_metadata.update(_standard_logging_metadata)
return _metadata
def _get_latency_metrics(self, standard_logging_payload: StandardLoggingPayload) -> DDLLMObsLatencyMetrics:
"""
Get the latency metrics from the standard logging payload
@ -808,7 +955,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0)
# Get budget information from metadata
metadata: Final = standard_logging_payload.get("metadata", {})
metadata: Final = _metadata_of(standard_logging_payload)
# API key max budget
user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget")

4
litellm/integrations/websearch_interception/handler.py
View file

@ -10,9 +10,9 @@ import asyncio
import math
import uuid
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast
from typing_extensions import ReadOnly
from typing_extensions import Never, ReadOnly
import litellm
from litellm._logging import verbose_logger

12
litellm/llms/a2a/chat/guardrail_translation/handler.py
View file

@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation):
return responses_so_far
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
_, valid_parsed = self._parse_streaming_responses(responses_so_far)
combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed)
return StreamingScanKey(texts=(combined_text,))
def _parse_streaming_responses(
self,
responses_so_far: list[object],
responses_so_far: Sequence[object],
) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]:
"""Parse JSON-RPC items, returning aligned parsed list and valid entries."""
parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far)

25
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
LiteLLMAnthropicMessagesAdapter,
is_provider_native_tool_dict,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
anthropic_tool_name,
anthropic_tool_names,
@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
merge_guardrailed_scoped_messages,
merge_returned_tools_into_request_tools,
scoped_structured_message_indices,
stream_item_fingerprint,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
stream_ended=stream_ended,
)
@classmethod
def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]:
return tuple(
stream_item_fingerprint(block)
for item in responses_so_far
for event in cls._iter_sse_events(item)
if event.get("type") == "content_block_start"
and isinstance(block := event.get("content_block"), Mapping)
and block.get("type") == "tool_use"
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Parse streaming responses and extract accumulated text content.

19
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -35,6 +35,22 @@ class StreamTransformSink:
holdback_per_choice: dict[int, int] = field(default_factory=dict)
@dataclass(frozen=True, slots=True)
class StreamingScanKey:
"""What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
compare equal when the round would scan the same content again; ``stream_ended``
stays out of the comparison and only says whether the handler is on its
end-of-stream path, where an empty payload is still scanned today."""
texts: tuple[str, ...]
tool_calls: tuple[str, ...] = ()
stream_ended: bool = field(default=False, compare=False)
@property
def has_nothing_to_scan(self) -> bool:
return not self.stream_ended and not any(self.texts) and not self.tool_calls
class BaseTranslation(ABC):
@staticmethod
def transform_user_api_key_dict_to_metadata(
@ -151,6 +167,9 @@ class BaseTranslation(ABC):
"""
return responses_so_far
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
return None
def build_block_sse_chunks(
self,
exc: "ModifyResponseException",

12
litellm/llms/base_llm/guardrail_translation/utils.py
View file

@ -4,6 +4,8 @@ import json
from collections.abc import Callable, Iterator, Sequence
from typing import Any, Final, TypeVar
from pydantic import BaseModel
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None:
return getattr(item, field, None)
def stream_item_fingerprint(item: object) -> str:
plain: Final = item.model_dump() if isinstance(item, BaseModel) else item
return json.dumps(plain, sort_keys=True, default=str)
def stream_item_items(item: object, field: str) -> tuple[object, ...]:
value: Final = stream_item_field(item, field)
return tuple(value) if isinstance(value, (list, tuple)) else ()
def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]:
"""
``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked

11
litellm/llms/fireworks_ai/common_utils.py
View file

@ -2,6 +2,7 @@ from typing import Final
from httpx import Headers
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
"""
Session id to send as `x-session-affinity`, or None when the caller gave none.
Deliberately does not fall back to `litellm_trace_id`: that is generated per
request (`str(uuid.uuid4())` when absent), so using it pins every request to a
different Fireworks node and prompt caching never hits.
Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the
proxy generated for a request that had none: both are per request, so using them
pins every request to a different Fireworks node and prompt caching never hits.
"""
params: Final = litellm_params
metadata: Final = params.get("metadata")
if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return None
for key in ("litellm_session_id", "session_id"):
value = params.get(key)
if value:
return str(value)
metadata: Final = params.get("metadata")
if isinstance(metadata, dict):
value = metadata.get("session_id")
if value:

37
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -26,6 +26,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
role_out_of_guardrail_scope,
scoped_structured_message_indices,
stream_item_field,
stream_item_fingerprint,
stream_item_items,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
terminate the stream. Text rewrites are not propagated to the client here
(see ``_process_streaming_transform`` for the incremental_diff path)."""
# check if the stream has ended
has_stream_ended = False
for chunk in responses_so_far:
if chunk.choices and chunk.choices[0].finish_reason is not None:
has_stream_ended = True
break
has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
if has_stream_ended:
# convert to model response
@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
}
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
stream_ended: Final = self._first_choice_has_finished(responses_so_far)
return StreamingScanKey(
texts=tuple(self._combine_streaming_texts(chunks).values()),
tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
stream_ended=stream_ended,
)
@staticmethod
def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]:
return tuple(
stream_item_fingerprint(tool_call)
for chunk in responses_so_far
for choice in _stream_chunk_choices(chunk)
for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls")
)
@staticmethod
def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool:
first_choices: Final = tuple(
choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices
)
return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices)
def _combine_streaming_texts(
self, responses_so_far: list["ModelResponseStream"]
self, responses_so_far: Sequence["ModelResponseStream"]
) -> dict[tuple[int, int | None], str]:
"""
Combine all streaming chunks into complete text per choice.

72
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
OpenAiResponsesToChatCompletionStreamIterator,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_stream_usage,
stream_item_field,
stream_item_fingerprint,
stream_item_items,
)
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
from litellm.responses.litellm_completion_transformation.transformation import (
@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation):
)
return responses_so_far
def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool:
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if the streaming has ended.
"""
if not responses_so_far:
return False
terminal_types: Final = {
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
}
return responses_so_far[-1].get("type") in terminal_types
terminal_types: Final = frozenset(
(
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
)
)
return stream_item_field(responses_so_far[-1], "type") in terminal_types
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
if not responses_so_far or not hasattr(responses_so_far[-1], "get"):
return None
last_event: Final = responses_so_far[-1]
last_event_type: Final = stream_item_field(last_event, "type")
if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
return None
if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
return self._completed_response_scan_key(stream_item_field(last_event, "response"))
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
stream_ended=self._check_streaming_has_ended(responses_so_far),
)
@staticmethod
def _completed_response_scan_key(response: object) -> StreamingScanKey:
output_items: Final = stream_item_items(response, "output")
message_items: Final = tuple(
item for item in output_items if stream_item_field(item, "type") != "function_call"
)
return StreamingScanKey(
texts=tuple(
text
for item in message_items
for part in stream_item_items(item, "content")
if isinstance(text := stream_item_field(part, "text"), str) and text
),
tool_calls=tuple(
stream_item_fingerprint(item)
for item in output_items
if stream_item_field(item, "type") == "function_call"
),
stream_ended=True,
)
def build_stream_error_items(
self,
@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation):
),
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Get the string so far from the responses so far.
@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation):
"""
keyed_events: Final = tuple(
(
(event.get("item_id"), event.get("output_index"), event.get("content_index")),
event.get("text"),
event.get("delta"),
(
stream_item_field(event, "item_id"),
stream_item_field(event, "output_index"),
stream_item_field(event, "content_index"),
),
stream_item_field(event, "text"),
stream_item_field(event, "delta"),
)
for event in responses_so_far
if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str)
)
def part_text(part_key: tuple[object, object, object]) -> str:

97
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -820,20 +820,46 @@ def _should_strip_caller_authorization(
if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate):
return False
normalized_raw_headers: Final = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)}
has_explicit_litellm_admission_header: Final = normalized_raw_headers.get("x-litellm-api-key") is not None
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
if mcp_server.is_oauth_delegate:
return not has_explicit_litellm_admission_header
admission_consumed_authorization_as_litellm_key: Final = (
user_api_key_auth is not None
and bool(getattr(user_api_key_auth, "api_key", None))
and not has_explicit_litellm_admission_header
)
return admission_consumed_authorization_as_litellm_key or (
return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or (
user_api_key_auth is None and not has_explicit_litellm_admission_header
)
LITELLM_VIRTUAL_KEY_PREFIX: Final = "sk-"
def _raw_header_value(raw_headers: Mapping[str, str] | None, name: str) -> str | None:
return next((v for k, v in (raw_headers or {}).items() if isinstance(k, str) and k.lower() == name), None)
def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None) -> bool:
"""Admission only consumes a non-empty ``x-litellm-api-key``; an empty one falls back to ``Authorization``."""
return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
def _authorization_is_litellm_admission_credential(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> bool:
"""True when ``Authorization`` carries the LiteLLM key admission validated.
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
same key in both headers.
"""
if user_api_key_auth is None or not user_api_key_auth.api_key:
return False
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
if not admission_header:
return True
authorization: Final = _raw_header_value(raw_headers, "authorization")
return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(
admission_header, "Bearer"
)
def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
"""Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection.
@ -3277,8 +3303,8 @@ class MCPServerManager:
#########################################################
@staticmethod
def _extract_bearer_token(
oauth2_headers: dict[str, str] | None,
raw_headers: dict[str, str] | None,
oauth2_headers: Mapping[str, str] | None,
raw_headers: Mapping[str, str] | None,
) -> str | None:
"""Extract the bare Bearer token from oauth2_headers or raw_headers.
@ -3298,10 +3324,29 @@ class MCPServerManager:
return auth_value
return None
@staticmethod
def _extract_subject_token(
oauth2_headers: Mapping[str, str] | None,
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""The caller's upstream identity token, or ``None`` when the bearer is a LiteLLM key.
Rejects the key admission validated and, because virtual keys always carry the ``sk-`` prefix,
any other LiteLLM key a client puts in ``Authorization`` next to ``x-litellm-api-key``.
"""
if _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth):
return None
bearer: Final = MCPServerManager._extract_bearer_token(oauth2_headers, raw_headers)
if bearer is not None and bearer.startswith(LITELLM_VIRTUAL_KEY_PREFIX):
return None
return bearer
def _obo_subject_token(
self,
server: MCPServer,
raw_headers: dict[str, str] | None,
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""The caller's bearer as the token_exchange (OBO) subject token, for that mode only.
@ -3311,7 +3356,7 @@ class MCPServerManager:
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
return None
return self._extract_bearer_token(None, raw_headers)
return self._extract_subject_token(None, raw_headers, user_api_key_auth)
def _build_stdio_env(
self,
@ -3566,6 +3611,7 @@ class MCPServerManager:
server: MCPServer,
oauth2_headers: dict[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Run the OBO exchange for a caller-supplied subject at the transport edge.
@ -3577,13 +3623,15 @@ class MCPServerManager:
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
return
subject_token: Final = self._extract_bearer_token(oauth2_headers, None)
if not subject_token:
if not self._extract_bearer_token(oauth2_headers, None):
return
resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
spec: Final = to_server_spec(resolved_server)
if spec is None or not isinstance(spec.config, TokenExchangeConfig):
return
subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
if subject_token is None:
raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(_):
return
@ -3851,7 +3899,7 @@ class MCPServerManager:
# token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes
# never read the inbound bearer, so leave subject_token None to avoid forwarding it.
subject_token: Final = (
self._extract_bearer_token(oauth2_headers, raw_headers)
self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
if server.auth_type == MCPAuth.oauth2_token_exchange
else None
)
@ -3931,6 +3979,7 @@ class MCPServerManager:
async def get_prompts_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -3959,7 +4008,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -3982,6 +4031,7 @@ class MCPServerManager:
async def get_resources_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -4001,7 +4051,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -4024,6 +4074,7 @@ class MCPServerManager:
async def get_resource_templates_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -4043,7 +4094,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -4068,6 +4119,7 @@ class MCPServerManager:
async def read_resource_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
url: AnyUrl,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
@ -4084,7 +4136,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client: Final = await self._create_mcp_client(
server=server,
@ -4099,6 +4151,7 @@ class MCPServerManager:
async def get_prompt_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
prompt_name: str,
arguments: dict[str, str] | None = None,
mcp_auth_header: str | dict[str, str] | None = None,
@ -4116,7 +4169,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client: Final = await self._create_mcp_client(
server=server,
@ -5290,7 +5343,7 @@ class MCPServerManager:
MCPAuth.oauth2_token_exchange,
MCPAuth.oauth2_id_jag,
):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
@ -5638,7 +5691,7 @@ class MCPServerManager:
subject_token: str | None = None
if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
elif isinstance(spec.config, PassthroughConfig):
inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers)
per_server_token: Final = _passthrough_token_from_mcp_auth_header(mcp_auth_header)

9
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -2189,6 +2189,7 @@ if MCP_AVAILABLE:
try:
prompts = await global_mcp_server_manager.get_prompts_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -2242,6 +2243,7 @@ if MCP_AVAILABLE:
try:
resources = await global_mcp_server_manager.get_resources_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -2293,6 +2295,7 @@ if MCP_AVAILABLE:
try:
resource_templates = await global_mcp_server_manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -3211,6 +3214,7 @@ if MCP_AVAILABLE:
return await global_mcp_server_manager.get_prompt_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
prompt_name=original_prompt_name,
arguments=arguments,
mcp_auth_header=server_auth_header,
@ -3261,6 +3265,7 @@ if MCP_AVAILABLE:
return await global_mcp_server_manager.read_resource_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
url=url,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
@ -3723,6 +3728,7 @@ if MCP_AVAILABLE:
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None,
allowed_server_ids: set[str] | None = None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Fail fast with HTTP 401 for MCP servers that need user auth but
didn't receive it on this request. Covers both gateway-managed OAuth2
@ -3867,6 +3873,7 @@ if MCP_AVAILABLE:
server=server,
oauth2_headers=oauth2_headers,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
)
# Pass-through OAuth: when the admin has opted a server into
@ -4195,6 +4202,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers. Must run after
@ -4518,6 +4526,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_sse_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers: surface upstream

4
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -5,10 +5,10 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never
from typing import TYPE_CHECKING, Any, Final, TypedDict
from pydantic import ValidationError
from typing_extensions import ReadOnly, Required
from typing_extensions import ReadOnly, Required, assert_never
import litellm
from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K

4
litellm/proxy/_types.py
View file

@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.",
)
missing_session_id: Literal["generate", "reject"] | None = Field(
None,
description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.",
)
enable_public_model_hub: bool = Field(
default=False,
description="Public model hub for users to see what models they have access to, supported openai params, etc.",

4
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -13,10 +13,10 @@ import os
import uuid
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final, TypedDict, assert_never
from typing import Annotated, Final, TypedDict
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from typing_extensions import ReadOnly, Required
from typing_extensions import ReadOnly, Required, assert_never
import litellm
from litellm._logging import verbose_proxy_logger

4
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -7,7 +7,9 @@ from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from types import MappingProxyType
from typing import Final, Literal, Protocol, TypeVar, assert_never
from typing import Final, Literal, Protocol, TypeVar
from typing_extensions import assert_never
import litellm
from litellm._logging import verbose_proxy_logger

64
litellm/proxy/guardrails/anthropic_sse.py
View file

@ -13,6 +13,19 @@ from typing import Final
from litellm.types.utils import Choices, ModelResponse
_ANTHROPIC_EVENT_TYPES: Final = frozenset(
{
"message_start",
"message_delta",
"message_stop",
"content_block_start",
"content_block_delta",
"content_block_stop",
"ping",
"error",
}
)
def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
@ -30,23 +43,43 @@ def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None:
return None
def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]:
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
return tuple(
event_data
for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses
if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing
)
def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
return next(
(
message
for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses
if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing
and event_data.get("type") == "message_start"
and isinstance(message := event_data.get("message"), dict)
for event_data in _parsed_sse_events(sse_stream)
if event_data.get("type") == "message_start" and isinstance(message := event_data.get("message"), dict)
),
None,
)
def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool:
"""Whether raw SSE frames are Anthropic Messages events.
``is_raw_sse_stream`` only says the chunks are unparsed bytes, and ``/v1/messages`` is not the
only endpoint that streams those: the Google ``:streamGenerateContent`` route marks its own
stream raw too. Reading its frames as Anthropic ones would refuse the response in a wire format
its client cannot parse, so the surface is decided on the event types actually present.
"""
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return False
return any(event.get("type") in _ANTHROPIC_EVENT_TYPES for event in _parsed_sse_events(sse_stream))
def assemble_anthropic_sse_stream(
all_chunks: Sequence[object], *, restore_identity: bool = False
) -> ModelResponse | None:
@ -111,6 +144,27 @@ def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]:
)
def is_sse_error_stream(all_chunks: Sequence[object]) -> bool:
"""Whether the buffered stream carries nothing but error frames.
post_call guardrails run in a chain, so a hook can be handed the terminal error frames an
earlier guardrail emitted when it blocked. Those carry no message to assemble, and replacing
them would hide the refusal the client is owed. Covers both wire forms a guardrail emits: the
Anthropic ``error`` event and the chat-completions ``{"error": ...}`` payload.
"""
if not all(isinstance(chunk, (str, bytes)) for chunk in all_chunks):
# A stream mixing typed chunks with an error frame still carries content to scan, and the
# frames-only join below would drop exactly the part that has to be scanned
return False
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return False
events: Final = _parsed_sse_events(sse_stream)
return len(events) > 0 and all(
event.get("type") == "error" or isinstance(event.get("error"), Mapping) for event in events
)
def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,

479
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -1,4 +1,5 @@
from collections.abc import AsyncGenerator, Mapping, Sequence
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
@ -27,12 +28,25 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.anthropic_sse import (
anthropic_sse_chunks_from_response,
anthropic_sse_error_frames,
assemble_anthropic_sse_stream,
is_anthropic_sse_stream,
is_raw_sse_stream,
is_sse_error_stream,
)
from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import (
MODEL_ARMOR_MAX_FILE_SIZE_BYTES,
plan_file_scans,
)
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolCallChunk,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
@ -41,10 +55,33 @@ from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
StandardLoggingGuardrailInformation,
TextCompletionResponse,
)
GUARDRAIL_NAME: Final = "model_armor"
# Only these carry the finished output; response.created carries an empty body
_RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
# Every event whose ``delta`` is model output already on its way to the client. Read off the event
# enum rather than listed, so an event added there cannot quietly fall out of the scan
_RESPONSES_DELTA_EVENT_TYPES: Final = frozenset(
event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta")
)
# What makes two delta events part of the same field of the turn, rather than two fields that merely
# streamed next to each other
_RESPONSES_DELTA_FIELD_ATTRS: Final = ("type", "item_id", "output_index", "content_index", "summary_index")
class _StreamSurface(Enum):
"""Wire format of a buffered streaming response, which decides how it is read and how it is refused."""
CHAT_COMPLETIONS = auto()
ANTHROPIC_MESSAGES = auto()
RESPONSES = auto()
OPAQUE_SSE = auto()
class ModelArmorAPIError(Exception):
"""Model Armor API failure (non-2xx), distinct from a content-block decision so
@ -322,19 +359,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
else:
return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> bool:
def _should_block_content(self, armor_response: Mapping[str, Any], allow_sanitization: bool = False) -> bool:
"""Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
sanitization_result: Final = armor_response.get("sanitizationResult", {})
filter_results: Final = sanitization_result.get("filterResults", {})
# filterResults can be a dict (named keys) or a list (array of filter result dicts)
filter_result_items = []
if isinstance(filter_results, dict):
filter_result_items = list(filter_results.values())
elif isinstance(filter_results, list):
filter_result_items = filter_results
for filt in filter_result_items:
for filt in self._filter_result_items(armor_response):
# Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before
if filt.get("raiFilterResult", {}).get("matchState") == "MATCH_FOUND":
return True
@ -358,22 +385,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
# Fallback dict code removed; all cases handled above
return False
def _get_sanitized_content(self, armor_response: dict) -> str | None:
def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None:
"""
Get the sanitized content from a Model Armor response, if available.
Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found.
"""
result: Final = armor_response.get("sanitizationResult", {})
filter_results: Final = result.get("filterResults", {})
# filterResults can be a dict (single filter) or a list (multiple filters)
filters: Final = (
list(filter_results.values())
if isinstance(filter_results, dict)
else filter_results
if isinstance(filter_results, list)
else []
)
filters: Final = self._filter_result_items(armor_response)
# Prefer sanitized text from deidentifyResult if present
for filter_entry in filters:
@ -397,6 +414,61 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
# Fallback: if Model Armor put sanitized text at the root, use it
return armor_response.get("sanitizedText") or armor_response.get("text")
@staticmethod
def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]:
"""Every filter result in a scan response.
filterResults is a dict of named filters on most templates and a list on some, so both
shapes are flattened to the same list of filter entries.
"""
filter_results: Final = armor_response.get("sanitizationResult", {}).get("filterResults", {})
if isinstance(filter_results, dict):
return list(filter_results.values())
if isinstance(filter_results, list):
return filter_results
return []
def _has_deidentify_match(self, armor_response: Mapping[str, Any]) -> bool:
"""Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction."""
for filter_entry in self._filter_result_items(armor_response):
sdp = filter_entry.get("sdpFilterResult")
if sdp and sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND":
return True
return False
def _resolve_streaming_outcome(
self,
armor_response: Mapping[str, Any],
assembled_response: object,
content: str,
) -> tuple[bool, str | None]:
"""Whether to block the buffered stream, and the rewrite to emit when it is not blocked.
A de-identify match only reaches here unblocked because masking is on, so the redaction it
stands for has to be both resolvable and emittable. Where it is neither, the buffered
original still carries what Model Armor matched on, so this fails closed instead of
releasing it.
"""
if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content):
return True, None
if not self.mask_response_content:
return False, None
sanitized_content: Final = self._get_sanitized_content(armor_response)
if not sanitized_content:
# No rewrite to apply. Harmless unless a match is outstanding, in which case applying
# nothing would hand back the very content that matched
return self._has_deidentify_match(armor_response), None
if sanitized_content == content:
return False, None
if not isinstance(assembled_response, ModelResponse):
verbose_proxy_logger.warning(
"Model Armor: sanitized content cannot be re-emitted on this streaming endpoint, "
"blocking the response instead"
)
return True, None
return False, sanitized_content
@staticmethod
def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object:
"""Accumulate scan responses so a later text scan does not drop an earlier file scan.
@ -831,6 +903,185 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return response
@staticmethod
def _is_terminal_error_stream(all_chunks: Sequence[object]) -> bool:
"""Whether the buffered stream is only the refusal an earlier guardrail in the chain emitted.
post_call guardrails are composed, so this hook can be handed the terminal error items a
preceding one produced. They carry no message to scan, and replacing them would hide the
refusal the client is owed.
"""
if all(getattr(chunk, "type", None) == "error" for chunk in all_chunks):
return True
return is_sse_error_stream(all_chunks)
@staticmethod
def _classify_stream(all_chunks: Sequence[object]) -> _StreamSurface:
"""Wire format the buffered chunks belong to."""
if is_raw_sse_stream(all_chunks):
return (
_StreamSurface.ANTHROPIC_MESSAGES if is_anthropic_sse_stream(all_chunks) else _StreamSurface.OPAQUE_SSE
)
if any(
isinstance(event_type := getattr(chunk, "type", None), str) and event_type.startswith("response.")
for chunk in all_chunks
):
return _StreamSurface.RESPONSES
return _StreamSurface.CHAT_COMPLETIONS
@staticmethod
def _final_responses_api_response(all_chunks: Sequence[object]) -> ResponsesAPIResponse | None:
"""Response body carried by a terminal ``/v1/responses`` event.
A stream cut short before it completes has to read as unassembled rather than as a clean
empty response: ``response.created`` also carries a body, but an empty one, and scanning
that would release every buffered delta unscanned.
"""
return next(
(
body
for chunk in reversed(all_chunks)
if getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES
and isinstance(body := getattr(chunk, "response", None), ResponsesAPIResponse)
),
None,
)
@staticmethod
def _responses_api_response_text(response: ResponsesAPIResponse) -> str:
"""Text to scan in a Responses API response, tool-call arguments included.
Tool calls are folded in because ``get_content_from_model_response`` folds them into what
the chat surface scans, and a Responses turn can carry its whole payload in them.
"""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists
tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink
handler: Final = OpenAIResponsesHandler()
for output_idx, output_item in enumerate(response.output or ()):
handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing
output_item=output_item,
output_idx=output_idx,
texts_to_check=texts,
images_to_check=[], # mutable-ok: the extractor's images sink, unused here
task_mappings=[], # mutable-ok: the extractor's task-mapping sink, unused here
tool_calls_to_check=tool_calls,
)
return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls)))
def _extract_streaming_content(self, assembled_response: object) -> str:
"""Text to scan from an assembled stream, for every endpoint shape this hook serves."""
if isinstance(assembled_response, ResponsesAPIResponse):
return self._responses_api_response_text(assembled_response)
return self._extract_content_from_response(assembled_response)
@staticmethod
def _responses_delta_field(chunk: object) -> tuple[str, ...]:
"""Which field of the turn a delta event belongs to."""
return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS)
@staticmethod
def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]:
"""Text each field of a ``/v1/responses`` turn has already spelled out in its delta events.
One field's deltas are joined as they streamed, since a finding can be split across them,
and separate fields stay apart, so a reasoning summary running into the visible answer
cannot spell out a finding that neither of them carries.
"""
deltas: Final = tuple(
(ModelArmorGuardrail._responses_delta_field(chunk), delta)
for chunk in all_chunks
if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
and isinstance(delta := getattr(chunk, "delta", None), str)
)
return tuple(
"".join(delta for field, delta in deltas if field == streamed_field)
for streamed_field in dict.fromkeys(field for field, _ in deltas)
)
def _streaming_content_to_scan(
self,
assembled_response: object,
all_chunks: Sequence[object],
surface: _StreamSurface,
) -> str:
"""Text to scan for a buffered stream, which is everything the client is about to receive.
A ``/v1/responses`` stream also spells out reasoning summaries and tool-call arguments in
delta events that its terminal body never repeats, so every delta field the body does not
already carry is scanned after it.
"""
content: Final = self._extract_streaming_content(assembled_response)
if surface is not _StreamSurface.RESPONSES:
return content
unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
return "\n".join(part for part in (content, *unscanned) if part)
@staticmethod
def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None:
"""Replace every non-empty choice message with the Model Armor sanitized text."""
for choice in assembled_response.choices:
if isinstance(choice, Choices) and choice.message.content:
choice.message.content = sanitized_content
@staticmethod
def _assemble_chat_completion_stream(
all_chunks: list[object], # mutable-ok: stream_chunk_builder only accepts a mutable list
) -> ModelResponse | TextCompletionResponse | None:
"""Assemble chat-completion chunks, returning ``None`` when they cannot be assembled."""
from litellm.main import stream_chunk_builder
try:
return stream_chunk_builder(chunks=all_chunks)
except Exception as exc:
verbose_proxy_logger.warning("Model Armor: chat-completion stream assembly failed (%s)", exc)
return None
def _assemble_stream(
self, all_chunks: Sequence[object], surface: _StreamSurface
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None:
"""Assemble the buffered stream into the scannable response its surface produces."""
if surface is _StreamSurface.ANTHROPIC_MESSAGES:
return assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
if surface is _StreamSurface.RESPONSES:
return self._final_responses_api_response(all_chunks)
if surface is _StreamSurface.OPAQUE_SSE:
return None
return self._assemble_chat_completion_stream(list(all_chunks))
@staticmethod
def _error_payload(exc: HTTPException) -> Mapping[str, object]:
"""Error object for a terminal stream item, carrying the status the frame would otherwise lose."""
detail: Final = exc.detail if isinstance(exc.detail, Mapping) else {"message": str(exc.detail)}
error_value: Final = detail.get("error", detail)
return {
**(dict(error_value) if isinstance(error_value, Mapping) else {"message": str(error_value)}),
"code": str(exc.status_code),
}
@staticmethod
def _build_responses_error_items(exc: HTTPException) -> Sequence[object] | None:
"""Responses API error events for a failure discovered after the stream started."""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
return OpenAIResponsesHandler().build_stream_error_items(exc, responses_so_far=None)
def _stream_error_items(self, exc: HTTPException, *, surface: _StreamSurface) -> Sequence[object]:
"""Frame a guardrail failure as terminal stream items in this endpoint's wire format."""
payload: Final = self._error_payload(exc)
if surface is _StreamSurface.ANTHROPIC_MESSAGES:
return anthropic_sse_error_frames(str(payload.get("message", "")))
if surface is _StreamSurface.RESPONSES and (responses_items := self._build_responses_error_items(exc)):
return responses_items
# Also the fallback when a surface cannot frame its own error: create_response() reads the
# status back out of this form, so the refusal keeps its code instead of arriving as a 200
return (f"data: {json.dumps({'error': payload})}\n\n",)
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -840,97 +1091,125 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
"""Process streaming response chunks."""
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
# Collect all chunks
all_chunks: Final[list[ModelResponseStream]] = []
all_chunks: Final[list[Any]] = []
async for chunk in response:
all_chunks.append(chunk)
if not all_chunks or self._is_terminal_error_stream(all_chunks):
for chunk in all_chunks:
yield chunk
return
surface: Final = self._classify_stream(all_chunks)
# Build complete response
assembled_response: Final = stream_chunk_builder(chunks=all_chunks)
assembled_response: Final = self._assemble_stream(all_chunks, surface)
if isinstance(assembled_response, ModelResponse):
# Extract content
content: Final = self._extract_content_from_response(assembled_response)
if assembled_response is None:
if not self.optional_params.get("fail_on_error", True):
verbose_proxy_logger.warning(
"Model Armor: streamed response could not be assembled for scanning, "
"forwarding it unscanned because fail_on_error is disabled"
)
for chunk in all_chunks:
yield chunk
return
if content:
try:
# Check with Model Armor
armor_response: Final = await self.make_model_armor_request(
content=content,
source="model_response",
request_data=request_data,
)
# Forwarding an unscannable stream would silently disable the guardrail, so fail closed
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
for error_item in self._stream_error_items(
HTTPException(
status_code=500,
detail=f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it",
),
surface=surface,
):
yield error_item
return
# Attach Model Armor response & status to this request's metadata to avoid race conditions
if isinstance(request_data, dict):
_, metadata = get_or_create_metadata_bucket(request_data)
metadata["_model_armor_response"] = self._build_logging_response(armor_response)
metadata["_model_armor_status"] = (
"blocked" if self._should_block_content(armor_response) else "success"
)
# Extract content
content: Final = self._streaming_content_to_scan(
assembled_response=assembled_response, all_chunks=all_chunks, surface=surface
)
# Add guardrail to applied_guardrails BEFORE potential blocking
# This ensures guardrail is recorded even when it blocks the request
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
if not content:
verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail")
for chunk in all_chunks:
yield chunk
return
add_guardrail_to_applied_guardrails_header(
request_data=request_data, guardrail_name=self.guardrail_name
)
try:
# Check with Model Armor
armor_response: Final = await self.make_model_armor_request(
content=content,
source="model_response",
request_data=request_data,
)
# Check if blocked
if self._should_block_content(armor_response):
raise HTTPException(
status_code=400,
detail=self._build_block_error_detail(
"Streaming response blocked by Model Armor",
armor_response,
),
)
# Decide the outcome before recording it. Mirrors the non-streaming sibling: with
# masking on, a de-identify match is a redaction to apply rather than a refusal, but
# that only holds while the redaction can actually be delivered
blocked, sanitized_content = self._resolve_streaming_outcome(
armor_response=armor_response,
assembled_response=assembled_response,
content=content,
)
# Apply sanitization if enabled
if self.mask_response_content:
sanitized_content: Final = self._get_sanitized_content(armor_response)
if sanitized_content and sanitized_content != content:
# Update assembled response
for choice in assembled_response.choices:
if isinstance(choice, Choices):
if choice.message.content:
choice.message.content = sanitized_content
# Attach Model Armor response & status to this request's metadata to avoid race conditions
if isinstance(request_data, dict):
_, metadata = get_or_create_metadata_bucket(request_data)
metadata["_model_armor_response"] = self._build_logging_response(armor_response)
metadata["_model_armor_status"] = "blocked" if blocked else "success"
# Return sanitized stream
mock_response: Final = MockResponseIterator(model_response=assembled_response)
async for chunk in mock_response:
yield chunk
return
# Add guardrail to applied_guardrails BEFORE potential blocking
# This ensures guardrail is recorded even when it blocks the request
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
except ModelArmorAPIError as e:
if self.optional_params.get("fail_on_error", True):
error_obj = {"message": e.detail, "code": "500"}
yield f"data: {json.dumps({'error': error_obj})}\n\n"
return
except HTTPException as e:
# Yield error as SSE event so create_response() detects it and
# returns a proper JSON error response with the correct status code.
# (Raising from a generator hits create_response's generic except → 500.)
detail: Final = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)}
error_value: Final = detail.get("error", detail)
if isinstance(error_value, dict):
error_obj = dict(error_value)
else:
error_obj = {"message": str(error_value)}
error_obj["code"] = str(e.status_code)
yield f"data: {json.dumps({'error': error_obj})}\n\n"
if blocked:
raise HTTPException(
status_code=400,
detail=self._build_block_error_detail(
"Streaming response blocked by Model Armor",
armor_response,
),
)
if sanitized_content is not None and isinstance(assembled_response, ModelResponse):
self._apply_sanitized_content(assembled_response, sanitized_content)
# Return sanitized stream
if surface is _StreamSurface.ANTHROPIC_MESSAGES:
for sse_chunk in anthropic_sse_chunks_from_response(assembled_response):
yield sse_chunk
return
except Exception as e:
verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True)
if self.optional_params.get("fail_on_error", True):
raise
else:
verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail")
mock_response: Final = MockResponseIterator(model_response=assembled_response)
async for chunk in mock_response:
yield chunk
return
except ModelArmorAPIError as e:
if self.optional_params.get("fail_on_error", True):
for error_item in self._stream_error_items(
HTTPException(status_code=500, detail=e.detail), surface=surface
):
yield error_item
return
except HTTPException as e:
# Yield the error as a terminal stream item so create_response() detects it and returns
# a proper JSON error response with the correct status code. Raising from a generator
# instead hits create_response's generic except and becomes a 500.
for error_item in self._stream_error_items(e, surface=surface):
yield error_item
return
except Exception as e:
verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True)
if self.optional_params.get("fail_on_error", True):
raise
# Return original chunks if no sanitization needed
for chunk in all_chunks:

40
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -36,6 +36,7 @@ if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
# Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error
@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol):
@property
def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ...
@property
def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ...
@property
def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]:
return choices
def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
if scan_key is None:
return False
return scan_key == last_scan_key or scan_key.has_nothing_to_scan
class _StreamTerminated(Exception):
"""Internal signal that the incremental transform stream has already emitted
its terminal chunks (block message or in-stream error) and must stop."""
@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# Drives how a block terminates the stream: continue the in-progress
# message (True) vs emit a standalone block message (False, buffered).
chunks_yielded = False
last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round
async for item in response:
chunk_counter += 1
@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger):
# Process chunk based on sampling rate
if chunk_counter % sampling_rate == 0:
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
if _is_redundant_scan(scan_key, last_scan_key):
verbose_proxy_logger.debug(
"Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
chunk_counter,
guardrail_to_apply.guardrail_name,
)
chunks_yielded = True
responses_yielded.append(item)
yield item
continue
verbose_proxy_logger.debug(
"Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
chunk_counter,
@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger):
# string, permanently losing this chunk's content.
original_item = copy.deepcopy(item)
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
try:
await endpoint_translation.process_output_streaming_response(
responses_so_far=responses_so_far,
@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger):
):
yield error_item
return
if scan_key is not None:
last_scan_key = scan_key
chunks_yielded = True
responses_yielded.append(original_item)
yield original_item
@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger):
# preserve the list, not clone every chunk (deepcopy would double
# peak memory for large responses).
buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None
end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
if _is_redundant_scan(end_scan_key, last_scan_key):
verbose_proxy_logger.debug(
"Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
guardrail_to_apply.guardrail_name,
)
for buffered_item in buffered_items or ():
yield buffered_item
for pending_item in pending_end_of_stream_items:
responses_yielded.append(pending_item)
yield pending_item
return
try:
await endpoint_translation.process_output_streaming_response(

55
litellm/proxy/litellm_pre_call_utils.py
View file

@ -16,6 +16,7 @@ from starlette.datastructures import Headers
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
@ -23,6 +24,7 @@ from litellm.constants import (
OTEL_SERVICE_NAME_METADATA_KEYS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
@ -40,6 +42,7 @@ from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LiteLLMRoutes,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
@ -47,6 +50,8 @@ from litellm.proxy._types import (
TeamCallbackMetadata,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_utils import get_request_route
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.callback_utils import (
decrypt_callback_vars,
get_metadata_variable_name_from_kwargs,
@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None:
return session_id
def _is_llm_inference_route(request: Request) -> bool:
route: Final = get_request_route(request)
return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
def apply_missing_session_id_policy(
data: dict[str, object], # mutable-ok: stamps session ids in place on the request body the pipeline threads through
_metadata_variable_name: str,
general_settings: Mapping[str, object] | None,
request: Request,
) -> None:
policy: Final = general_settings.get("missing_session_id") if general_settings else None
if policy is None or not _is_llm_inference_route(request):
return
metadata: Final = data.get(_metadata_variable_name)
if not isinstance(metadata, dict):
return
if data.get("litellm_session_id") or metadata.get("session_id"):
return
match policy:
case "generate":
session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4())
data["litellm_session_id"] = session_id # rebind-ok: data is an out-param
data.setdefault("litellm_trace_id", session_id)
metadata["session_id"] = session_id
metadata[SESSION_ID_GENERATED_METADATA_KEY] = True
case "reject":
raise ProxyException(
message=(
"Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. "
"Required by `general_settings.missing_session_id: reject`."
),
type=ProxyErrorTypes.bad_request_error,
param="session_id",
code=400,
)
case _:
verbose_proxy_logger.warning(
"Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy
)
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
extensions and the Agent SDK run through the same CLI and share that prefix."""
@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request(
data=data,
_metadata_variable_name=_metadata_variable_name,
)
apply_missing_session_id_policy(
data=data,
_metadata_variable_name=_metadata_variable_name,
general_settings=general_settings,
request=request,
)
# Expose request headers under the metadata field for guardrails (fixes #17477)
if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict):

4
litellm/proxy/openai_files_endpoints/batch_file_validation.py
View file

@ -2,7 +2,9 @@ import json
from collections.abc import Iterator
from dataclasses import dataclass
from itertools import chain
from typing import BinaryIO, Final, NoReturn, assert_never
from typing import BinaryIO, Final, NoReturn
from typing_extensions import assert_never
from litellm.proxy._types import ProxyException

4
litellm/proxy/openai_files_endpoints/general_upload_validation.py
View file

@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose.
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Final, NoReturn, assert_never
from typing import BinaryIO, Final, NoReturn
from typing_extensions import assert_never
from litellm.proxy._types import ProxyException
from litellm.proxy.common_utils.path_utils import safe_filename

4
litellm/proxy/public_endpoints/autorouter_presets.json
View file

@ -17,6 +17,7 @@
"classification_mode": "every_request",
"session_affinity": false,
"modality_routing": false,
"modality_pin_override": false,
"deployment_affinity": true
}
},
@ -35,6 +36,7 @@
"classification_mode": "every_request",
"session_affinity": false,
"modality_routing": false,
"modality_pin_override": false,
"deployment_affinity": true
}
},
@ -63,6 +65,7 @@
"classification_mode": "every_request",
"session_affinity": false,
"modality_routing": false,
"modality_pin_override": false,
"deployment_affinity": true
}
},
@ -84,6 +87,7 @@
"classification_mode": "every_request",
"session_affinity": false,
"modality_routing": false,
"modality_pin_override": false,
"deployment_affinity": true
}
}

30
litellm/router.py
View file

@ -150,8 +150,10 @@ from litellm.router_utils.fallback_event_handlers import (
_check_non_standard_fallback_format,
clear_pre_routing_selection,
fallback_lookup_groups,
fallbacks_disabled_for_request,
get_fallback_model_group_for_lookup_groups,
get_pre_routing_selection,
record_disable_fallbacks,
record_pre_routing_selection,
run_async_fallback,
)
@ -5193,7 +5195,7 @@ class Router:
if not has_generated_content and error_event is None
else None
)
if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs):
if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs):
refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
raise MidStreamFallbackError(
message=refusal_error.message,
@ -7266,6 +7268,7 @@ class Router:
_fallback_metadata["original_model_group"] = model_group
include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True
disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False)
record_disable_fallbacks(kwargs, disable_fallbacks is True)
fallbacks: Final[list | None] = kwargs.get("fallbacks", self.fallbacks)
context_window_fallbacks: list | None = kwargs.get("context_window_fallbacks", self.context_window_fallbacks)
content_policy_fallbacks: list | None = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
@ -8131,6 +8134,29 @@ class Router:
)
return False
def _refusal_fallback_available(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
"""
Whether a safeguard refusal can actually be recovered by the dispatcher. A configured
content-policy list is authoritative; with none configured at all, the dispatcher falls
through to the generic fallbacks lookup, so the gate mirrors that reachability and arms
on a resolving generic chain (tier first, then the requested group, then "*").
"""
if fallbacks_disabled_for_request(kwargs):
return False
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
if content_policy_fallbacks is not None:
return self._has_content_policy_fallback(model_group, kwargs)
if self._has_default_fallbacks():
return True
fallbacks: Final = kwargs.get("fallbacks", self.fallbacks)
if fallbacks is None:
return False
resolved, _ = get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks,
lookup_groups=fallback_lookup_groups(kwargs, model_group),
)
return resolved is not None
def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool:
"""
Determines if a content policy error should be raised.
@ -8162,7 +8188,7 @@ class Router:
return False
if get_safeguard_refusal_stop_details(response) is None:
return False
return self._has_content_policy_fallback(model, kwargs)
return self._refusal_fallback_available(model, kwargs)
def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None):
_all_deployments: list = []

12
litellm/router_strategy/complexity_router/README.md
View file

@ -187,6 +187,9 @@ model_list:
# Replace a routed model that cannot take image input (default: false)
modality_routing: true
# Let that replacement also override a kept session pin, for image turns only (default: false)
modality_pin_override: true
```
## Usage
@ -227,9 +230,16 @@ vision model sits below the decided tier gets the 400 and an actionable message
A same-tier re-pick keeps the decision's cause and adds `modality:image` to `signals`; a tier
change or default takeover records `cause: modality_escalation` with the displaced placement
(`modality_escalated_from:<TIER>` or `modality_displaced_default_model`). Escalations are never
pinned by session affinity, and a KEPT session pin bypasses the gate entirely: a session pinned
pinned by session affinity, and by default a KEPT session pin bypasses the gate: a session pinned
to a text-only model keeps it even when an image arrives.
Add `modality_pin_override: true` to lift that last exemption. The image turn is then re-placed
the same way every other decision is, and records `cause: modality_pin_override` whether or not
the tier moved, since the model left the pin either way. The pin itself is untouched: the session
affinity write happens upstream of the gate and stores the session's own model, so the next text
turn replays the original pin and the override is never pinned in its place. It does nothing
unless `modality_routing` is also on.
### Heuristic-first chaining
`classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM

31
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
from pydantic import BaseModel, create_model
from litellm._logging import verbose_router_logger
from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY
from litellm.constants import (
EMPTY_MAPPING,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
@ -747,7 +751,8 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
A modality escalation is transient the same way: it describes what this one call carries (an
image), not what the session's traffic looks like, and pinning it would hold every following
text turn on the vision-capable model the image forced.
text turn on the vision-capable model the image forced. A modality pin override is the same
fact on a session that already holds a pin, so it must not overwrite the pin it displaced.
"""
return decision is None or (
decision.get("cause")
@ -756,6 +761,7 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
"plan_mode",
"housekeeping",
"modality_escalation",
"modality_pin_override",
)
and not decision.get("context_escalated")
)
@ -2389,8 +2395,11 @@ class ComplexityRouter(CustomLogger):
"""Replace a routed model that cannot accept this request's image input.
The single modality owner, applied to the decided response at the hook's exits so every
routing path is covered uniformly. A KEPT session pin is exempt by design (its cause);
replacement picks and every other path are just responses. The re-placement walks
routing path is covered uniformly. A KEPT session pin is exempt by design (its cause)
unless modality_pin_override is set, in which case the image turn is re-placed and reported
as modality_pin_override while the stored pin, written upstream from the session's own
model, is left for the next text turn; replacement picks and every other path are just
responses. The re-placement walks
UPWARD-ONLY from the decision's tier (so a plan-mode floor can never be undercut), picks
through `_pick_model_for_tier` so routing plugins still apply, then falls to
default_model (never on plugin routers, and never on a plan-floored decision, since
@ -2403,7 +2412,11 @@ class ComplexityRouter(CustomLogger):
not self.config.modality_routing
or not resolved_messages
or response.model is None
or (decision is not None and decision.get("cause") == "session_affinity_pin")
or (
decision is not None
and decision.get("cause") == "session_affinity_pin"
and not self.config.modality_pin_override
)
or not request_contains_image_content(resolved_messages)
or self._model_accepts_image_input(response.model)
):
@ -2445,6 +2458,10 @@ class ComplexityRouter(CustomLogger):
self._restamp_adaptive_choice(request_kwargs, response.model, new_model)
same_tier: Final = capable is not None and decided == capable
base_cause: Final = (decision.get("cause") if decision is not None else None) or "default_fallback"
# Reaching here on a kept pin means modality_pin_override is on, since the guard above
# returns otherwise. The model moved off the pin even on a same-tier repick, so reporting
# the pin's own cause would claim the session's model served a request it did not.
displaced_pin: Final = base_cause == "session_affinity_pin"
displaced_default: Final = decided is None and response.model == self.config.default_model
markers: Final = (
"modality:image",
@ -2454,7 +2471,7 @@ class ComplexityRouter(CustomLogger):
old_signals: Final = tuple(decision.get("signals") or ()) if decision is not None else ()
new_decision: Final = self._build_routing_decision(
routed_model=new_model,
cause=base_cause if same_tier else "modality_escalation",
cause="modality_pin_override" if displaced_pin else (base_cause if same_tier else "modality_escalation"),
tier=new_tier,
score=decision.get("score") if decision is not None else None,
signals=(*old_signals, *markers),
@ -2712,7 +2729,7 @@ class ComplexityRouter(CustomLogger):
"""Resolve a client-supplied session_id."""
for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
session_id = metadata.get("session_id")
if session_id is not None:
if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return str(session_id)
return None

15
litellm/router_strategy/complexity_router/config.py
View file

@ -883,7 +883,20 @@ class ComplexityRouterConfig(BaseModel):
"a routed model explicitly declared supports_vision false (deployment model_info "
"or the model cost map; unmapped names stay routable) is replaced by the nearest "
"HIGHER tier holding a capable model, then default_model, else a clear 400. A kept "
"session-affinity pin still wins even when an image arrives."
"session-affinity pin still wins even when an image arrives, unless "
"modality_pin_override is also enabled."
),
)
modality_pin_override: bool = Field(
default=False,
description=(
"Let modality_routing replace a kept session-affinity pin on the turns that carry an "
"image. Without this, a session pinned to a text-only model fails every image turn with "
"a provider 400, since the pin is exempt from the modality gate. When enabled, such a "
"turn routes to a capable model for that request only and the stored pin is left "
"untouched, so the next text turn replays the session's own model; the override is "
"reported as cause modality_pin_override and is never itself pinned. Inert unless "
"modality_routing is also enabled."
),
)

32
litellm/router_utils/fallback_event_handlers.py
View file

@ -263,6 +263,38 @@ def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None:
return next((selected for selected in selections if isinstance(selected, str) and selected), None)
DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks"
def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bool) -> None:
"""
Write-or-clear the request's disable_fallbacks verdict into the router-internal metadata
bucket. The wrapper pops the raw kwarg before any downstream frame runs, so the refusal
gate (which decides whether to convert a refusal into a recoverable error) needs this
carrier to know recovery is impossible.
"""
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
if request_kwargs is None:
return
bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))
if not isinstance(bucket, dict):
return
if disabled:
bucket[DISABLE_FALLBACKS_METADATA_KEY] = True
else:
bucket.pop(DISABLE_FALLBACKS_METADATA_KEY, None)
def fallbacks_disabled_for_request(kwargs: Mapping[str, Any]) -> bool:
"""True when this request opted out of fallbacks, read from the raw kwarg (pre-pop
snapshots keep it) or the router-internal bucket the wrapper stamps after popping it."""
if kwargs.get("disable_fallbacks") is True:
return True
buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS)
return any(isinstance(bucket, dict) and bucket.get(DISABLE_FALLBACKS_METADATA_KEY) is True for bucket in buckets)
def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]:
"""
Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins,

4
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

@ -21,7 +21,7 @@ from typing_extensions import TypedDict
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import AllMessageValues
@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger):
@staticmethod
def _get_session_id_from_metadata_dict(metadata: dict) -> str | None:
session_id: Final = metadata.get("session_id")
if session_id is None:
if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return None
return str(session_id)

1
litellm/types/integrations/datadog_llm_obs.py
View file

@ -86,6 +86,7 @@ class LLMMetrics(TypedDict, total=False):
cache_read_input_tokens: ReadOnly[float]
cache_write_input_tokens: ReadOnly[float]
non_cached_input_tokens: ReadOnly[float]
reasoning_output_tokens: ReadOnly[float]
class LLMObsPayload(TypedDict, total=False):

4
litellm/types/llms/gemini_audio_transcription.py
View file

@ -1,7 +1,7 @@
from typing import Literal, Required
from typing import Literal
from pydantic import BaseModel, ConfigDict
from typing_extensions import ReadOnly, TypedDict
from typing_extensions import ReadOnly, Required, TypedDict
class GeminiTranscriptionAudioInput(TypedDict):

4
litellm/types/utils.py
View file

@ -2879,6 +2879,10 @@ RoutingDecisionCause = Literal[
# routed model does not accept image input, so the nearest higher capable tier or
# default_model served instead. The displaced placement rides in signals.
"modality_escalation",
# modality_pin_override replaced a KEPT session-affinity pin for this request only: the turn
# carries an image the pinned model cannot accept. The stored pin is untouched, so the next
# text turn replays it. Distinct from "modality_escalation", which never displaces a pin.
"modality_pin_override",
"session_affinity_pin",
"session_affinity_escalation",
# classification_mode 'user_turn': the request is an agent loop's continuation turn (no new

150
tests/code_coverage_tests/check_py310_typing_imports.py Normal file
View file

@ -0,0 +1,150 @@
import ast
import os
import sys
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset(
{
"NotRequired",
"Required",
"Self",
"LiteralString",
"Never",
"assert_never",
"assert_type",
"reveal_type",
"TypeVarTuple",
"Unpack",
"dataclass_transform",
"override",
"TypeAliasType",
"get_original_bases",
"ReadOnly",
"TypeIs",
"NoDefault",
"get_protocol_members",
"is_protocol",
"evaluate_forward_ref",
"TypeForm",
}
)
@dataclass(frozen=True, slots=True)
class TypingImportViolation:
file: str
line: int
name: str
def _walk_with_ancestors(
node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = ()
) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]:
yield node, ancestors
for field_name, field_value in ast.iter_fields(node):
if isinstance(field_value, ast.AST):
yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name)))
elif isinstance(field_value, list):
for child in field_value:
if isinstance(child, ast.AST):
yield from _walk_with_ancestors(child, (*ancestors, (node, field_name)))
def _is_sys_version_info(node: ast.AST) -> bool:
return (
isinstance(node, ast.Attribute)
and isinstance(node.value, ast.Name)
and node.value.id == "sys"
and node.attr == "version_info"
)
def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool:
nearest_if: Final[tuple[ast.If, str] | None] = next(
(
(ancestor, field_name)
for ancestor, field_name in reversed(ancestors)
if isinstance(ancestor, ast.If)
),
None,
)
if nearest_if is None:
return False
enclosing_if, branch = nearest_if
test: Final[ast.expr] = enclosing_if.test
if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left):
return False
operator: Final[ast.cmpop] = test.ops[0]
return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or (
isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse"
)
def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]:
path: Final[Path] = Path(file_path)
tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
return tuple(
violation
for node, ancestors in _walk_with_ancestors(tree)
if not _is_version_guarded(ancestors)
for violation in _violations_for_node(node, path)
)
def _violations_for_node(
node: ast.AST, path: Path
) -> tuple[TypingImportViolation, ...]:
if isinstance(node, ast.ImportFrom) and node.module == "typing":
return tuple(
TypingImportViolation(file=str(path), line=node.lineno, name=alias.name)
for alias in node.names
if alias.name in PY311_PLUS_TYPING_NAMES
)
if (
isinstance(node, ast.Attribute)
and isinstance(node.value, ast.Name)
and node.value.id == "typing"
and node.attr in PY311_PLUS_TYPING_NAMES
):
return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),)
return ()
def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]:
base_path: Final[Path] = Path(base_dir)
return tuple(
violation
for directory in (
base_path / "litellm",
base_path / "enterprise",
base_path / "litellm-proxy-extras" / "litellm_proxy_extras",
)
if directory.exists()
for path in directory.rglob("*.py")
for violation in scan_file(path)
)
def main() -> None:
violations: Final[tuple[TypingImportViolation, ...]] = scan_directory()
if violations:
message: Final[str] = "\n".join(
(
"Python 3.10-incompatible typing imports found:",
*(
f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; "
"import it from typing_extensions instead because litellm supports Python 3.10"
for violation in violations
),
)
)
sys.stdout.write(f"{message}\n")
raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead")
sys.stdout.write("No Python 3.10-incompatible typing imports found.\n")
if __name__ == "__main__":
main()

7
tests/openai_endpoints_tests/test_e2e_openai_responses_api.py
View file

@ -1,5 +1,5 @@
import httpx
from openai import OpenAI, BadRequestError, APIStatusError
from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError
import pytest
@ -105,10 +105,9 @@ def test_streaming_response():
assert len(collected_chunks) > 0
def test_bad_request_error():
def test_model_not_found_error():
client = get_test_client()
with pytest.raises(BadRequestError):
# Trigger error with invalid model name
with pytest.raises(NotFoundError):
client.responses.create(model="non-existent-model", input="This should fail")

106
tests/router_unit_tests/test_router_anthropic_messages_fallback.py
View file

@ -338,6 +338,112 @@ def test_record_pre_routing_selection_writes_only_the_internal_bucket():
assert kwargs["metadata"] == {"user_id": "u1"}
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
async def test_generic_only_row_recovers_safeguard_refusal(stream):
"""With no content-policy list configured, a generic fallback row covers safeguard refusals,
so the dashboard's generic fallbacks work without config-only content_policy rows."""
fake = FakeAnthropicUpstream()
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
with fake.install():
response = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, stream=stream, messages=[{"role": "user", "content": "hi"}]
)
body = await _collect(response) if stream else response
if stream:
assert b'"refusal"' not in body
assert b"text_delta" in body
else:
assert body["stop_reason"] == "end_turn"
assert len(fake.calls) == 2
assert "claude-opus-5" in fake.calls[1]
@pytest.mark.asyncio
async def test_configured_content_policy_list_stays_authoritative_over_generic_rows():
fake = FakeAnthropicUpstream()
router = Router(
model_list=[FABLE_TIER, OPUS_TARGET],
fallbacks=[{"fable-tier": ["opus-target"]}],
content_policy_fallbacks=[{"unrelated-group": ["opus-target"]}],
)
with fake.install():
response = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}]
)
assert response["stop_reason"] == "refusal"
assert len(fake.calls) == 1
def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy():
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}])
stamped = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}}
assert router._refusal_fallback_available("router-group", stamped) is True
assert router._refusal_fallback_available("router-group", {}) is False
assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
def test_chat_content_filter_gate_unchanged_by_generic_rows():
"""The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's
content_filter gate keeps its long-standing content-policy-only semantics."""
from litellm.types.utils import Choices, ModelResponse
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
response = ModelResponse(choices=[Choices(finish_reason="content_filter")])
assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
async def test_disable_fallbacks_returns_the_refusal_instead_of_raising(stream):
"""A request that opted out of fallbacks must receive the provider's refusal response,
never a ContentPolicyViolationError the dispatcher refuses to recover."""
fake = FakeAnthropicUpstream()
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
with fake.install():
response = await router.aanthropic_messages(
model="fable-tier",
max_tokens=16,
stream=stream,
disable_fallbacks=True,
messages=[{"role": "user", "content": "hi"}],
)
body = await _collect(response) if stream else response
if stream:
assert b'"stop_reason": "refusal"' in body
else:
assert body["stop_reason"] == "refusal"
assert len(fake.calls) == 1
@pytest.mark.asyncio
async def test_disable_fallbacks_beats_a_content_policy_row_too():
fake = FakeAnthropicUpstream()
router = Router(
model_list=[FABLE_TIER, OPUS_TARGET],
content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
)
with fake.install():
response = await router.aanthropic_messages(
model="fable-tier",
max_tokens=16,
disable_fallbacks=True,
messages=[{"role": "user", "content": "hi"}],
)
assert response["stop_reason"] == "refusal"
assert len(fake.calls) == 1
def test_refusal_gate_keys_on_pre_routing_tier_stamp():
router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}])

9
tests/rust-python-harness/README.md
View file

@ -8,14 +8,17 @@ The matrix always has these SDK columns:
- `messages / amessages`
- `responses / aresponses`
- `count_tokens`
- `chat_completions / acompletion`
- `transcription / atranscription`
The harness has three deliberately broad test-strategy folders:
The harness has four deliberately broad test-strategy folders:
| Strategy | Folder |
| --- | --- |
| Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) |
| Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) |
| Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) |
| Already-existing live-API SDK tests | [`existing_e2e_test_sdk/`](existing_e2e_test_sdk/) |
## Run it
@ -112,7 +115,7 @@ The initial end-to-end entries deliberately show `◐`: the repository has Rust
## Attach parity tests
Each of the three folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
Each of the four folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
```json
{
@ -123,7 +126,7 @@ Each of the three folders contains a concise `README.md` and a `strategy.json`.
}
```
Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family; a selector ending in `/` aggregates every test in that folder, recursively. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
Use these coverage values:

49
tests/rust-python-harness/cli.py
View file

@ -6,9 +6,10 @@ from collections.abc import Sequence
from pathlib import Path
from .catalog import load_catalog
from .models import HarnessCase, Strategy
from .models import SDK_FUNCTIONS, HarnessCase, Strategy
from .runner import run_pytest
from .ui import make_dashboard
from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report
REPO_ROOT = Path(__file__).resolve().parents[2]
COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness"
@ -40,9 +41,17 @@ def _parser() -> argparse.ArgumentParser:
action="append",
default=[],
dest="sdk_functions",
choices=("ocr", "messages", "responses", "count_tokens"),
choices=SDK_FUNCTIONS,
help="run only this SDK function",
)
parser.add_argument(
"--validate-ledger",
action="store_true",
help=(
"report Python<->Rust test-parity ledger gaps and drift instead of "
"running the dashboard; narrow with --function"
),
)
parser.add_argument(
"--plain",
action="store_true",
@ -100,7 +109,7 @@ def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[
)
sdk_functions = _pick_values(
"SDK functions",
[(name, name) for name in ("ocr", "messages", "responses", "count_tokens")],
[(name, name) for name in SDK_FUNCTIONS],
)
return strategy_ids, sdk_functions
@ -131,6 +140,38 @@ def _print_catalog(strategies: Sequence[Strategy]) -> None:
print(f" {case.sdk_function:12} {case.coverage.value:14} {selectors}")
def _print_function_report(report: FunctionReport) -> None:
print(f"\n{report.sdk_function}")
if report.ledger is None or report.audit is None:
print(" no ledger yet")
return
ledger, audit = report.ledger, report.audit
print(
f" {ledger.mapped_count}/{ledger.total_count} python tests mapped to rust "
f"({ledger.percentage}%)"
)
print(f" {len(ledger.rust_only_tests)} rust-only tests with no python counterpart")
if audit.is_clean:
print(" ledger is in sync with the live test files")
return
for label, items in (
("ledger references a python test that no longer exists", audit.missing_python_tests),
("python test exists but is not tracked in the ledger", audit.stale_python_tests),
("ledger references a rust test that no longer exists", audit.missing_rust_tests),
("rust test exists but is not tracked in the ledger", audit.stale_rust_tests),
):
for item in items:
print(f" {label}: {item}")
def _validate_ledger(sdk_functions: set[str]) -> int:
functions = sdk_functions or set(SDK_FUNCTIONS)
reports = tuple(build_function_report(function) for function in sorted(functions))
for report in reports:
_print_function_report(report)
return 0 if all(report.is_clean for report in reports) else 1
def main(argv: Sequence[str] | None = None) -> int:
args = _parser().parse_args(argv)
if args.coverage and importlib.util.find_spec("pytest_cov") is None:
@ -138,6 +179,8 @@ def main(argv: Sequence[str] | None = None) -> int:
"--coverage requires the project's pytest-cov dependency; run with "
"`poetry run python -m tests.rust-python-harness --coverage`"
)
if args.validate_ledger:
return _validate_ledger(set(args.sdk_functions))
strategies = load_catalog()
if args.list:
_print_catalog(strategies)

4
tests/rust-python-harness/e2e_fuzz_tests/strategy.json
View file

@ -7,6 +7,8 @@
"ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"messages": {"coverage": "partial", "selectors": ["tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"responses": {"coverage": "partial", "selectors": ["tests/test_litellm/responses/test_rust_bridge_websocket.py"], "note": "Covers the websocket bridge; full responses parity is still being added."},
"count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."}
"count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."},
"chat_completions": {"coverage": "partial", "selectors": ["tests/test_litellm/rust_bridge/test_chat_completions.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"transcription": {"coverage": "partial", "selectors": ["tests/test_litellm/test_audio_transcription_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}
}
}

3
tests/rust-python-harness/existing_e2e_test_sdk/README.md Normal file
View file

@ -0,0 +1,3 @@
# Existing e2e SDK tests
Wires already-existing live-API SDK tests into the matrix instead of writing new parity tests. Selectors point at real test files and folders, such as `tests/ocr_tests/`, rather than individual node IDs, so future tests added to those folders are picked up automatically.

14
tests/rust-python-harness/existing_e2e_test_sdk/strategy.json Normal file
View file

@ -0,0 +1,14 @@
{
"order": 40,
"id": "existing_e2e_test_sdk",
"label": "Existing e2e SDK tests",
"description": "Wire already-existing live-API SDK tests into the matrix instead of writing new parity tests.",
"functions": {
"ocr": {"coverage": "partial", "selectors": ["tests/ocr_tests/"], "note": "Existing live OCR provider tests; not yet a frozen Rust/Python oracle comparison."},
"messages": {"coverage": "planned", "selectors": []},
"responses": {"coverage": "planned", "selectors": []},
"count_tokens": {"coverage": "planned", "selectors": []},
"chat_completions": {"coverage": "partial", "selectors": ["tests/llm_translation/test_anthropic_completion.py", "tests/llm_translation/test_bedrock_completion.py"], "note": "Existing live chat completion tests for providers with confirmed Rust bridge regressions."},
"transcription": {"coverage": "partial", "selectors": ["tests/audio_tests/test_whisper.py"], "note": "Existing live Whisper transcription test."}
}
}

2
tests/rust-python-harness/models.py
View file

@ -33,7 +33,7 @@ class ConfidenceLevel(str, Enum):
LOW = "LOW"
SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens")
SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription")
@dataclass(frozen=True)

2
tests/rust-python-harness/runner.py
View file

@ -15,6 +15,8 @@ UpdateCallback = Callable[[HarnessRun], None]
def selector_matches_node(selector: str, nodeid: str) -> bool:
normalized_selector = selector.replace("\\", "/")
normalized_nodeid = nodeid.replace("\\", "/")
if normalized_selector.endswith("/"):
return normalized_nodeid.startswith(normalized_selector)
if "::" in normalized_selector:
return normalized_nodeid == normalized_selector or normalized_nodeid.startswith(
f"{normalized_selector}["

0
tests/rust-python-harness/shared/__init__.py Normal file
View file

0
tests/rust-python-harness/shared/parity/__init__.py Normal file
View file

136
tests/rust-python-harness/shared/parity/ledger.py Normal file
View file

@ -0,0 +1,136 @@
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any
@dataclass(frozen=True, slots=True)
class LedgerEntry:
python_file: str
python_test: str
status: str
rust_file: str
rust_test: str
justification: str
reason: str
@dataclass(frozen=True, slots=True)
class RustOnlyEntry:
rust_file: str
rust_test: str
reason: str
@dataclass(frozen=True, slots=True)
class TestLedger:
sdk_function: str
python_scope: tuple[str, ...]
rust_scope: tuple[str, ...]
entries: tuple[LedgerEntry, ...]
rust_only_tests: tuple[RustOnlyEntry, ...]
@property
def mapped_count(self) -> int:
return sum(1 for entry in self.entries if entry.status == "mapped")
@property
def total_count(self) -> int:
return len(self.entries)
@property
def percentage(self) -> float:
if self.total_count == 0:
return 0.0
return round(100.0 * self.mapped_count / self.total_count, 1)
def _require_string(value: Any, field: str, source: Path) -> str:
if not isinstance(value, str) or not value.strip():
raise ValueError(f"{source}: {field} must be a non-empty string")
return value
def _require_string_list(value: Any, field: str, source: Path) -> tuple[str, ...]:
if not isinstance(value, list) or not all(isinstance(item, str) and item for item in value):
raise ValueError(f"{source}: {field} must be a list of non-empty strings")
return tuple(value)
def _load_entry(data: Any, index: int, source: Path) -> LedgerEntry:
if not isinstance(data, dict):
raise ValueError(f"{source}: entries[{index}] must be an object")
python_file = _require_string(data.get("python_file"), f"entries[{index}].python_file", source)
python_test = _require_string(data.get("python_test"), f"entries[{index}].python_test", source)
status = data.get("status")
if status not in ("mapped", "unmapped"):
raise ValueError(f"{source}: entries[{index}].status must be 'mapped' or 'unmapped'")
if status == "mapped":
rust_file = _require_string(data.get("rust_file"), f"entries[{index}].rust_file", source)
rust_test = _require_string(data.get("rust_test"), f"entries[{index}].rust_test", source)
justification = _require_string(
data.get("justification"), f"entries[{index}].justification", source
)
return LedgerEntry(
python_file=python_file,
python_test=python_test,
status=status,
rust_file=rust_file,
rust_test=rust_test,
justification=justification,
reason="",
)
reason = _require_string(data.get("reason"), f"entries[{index}].reason", source)
return LedgerEntry(
python_file=python_file,
python_test=python_test,
status=status,
rust_file="",
rust_test="",
justification="",
reason=reason,
)
def _load_rust_only_entry(data: Any, index: int, source: Path) -> RustOnlyEntry:
if not isinstance(data, dict):
raise ValueError(f"{source}: rust_only_tests[{index}] must be an object")
return RustOnlyEntry(
rust_file=_require_string(data.get("rust_file"), f"rust_only_tests[{index}].rust_file", source),
rust_test=_require_string(data.get("rust_test"), f"rust_only_tests[{index}].rust_test", source),
reason=_require_string(data.get("reason"), f"rust_only_tests[{index}].reason", source),
)
def load_ledger(path: Path) -> TestLedger:
with path.open(encoding="utf-8") as stream:
data = json.load(stream)
sdk_function = _require_string(data.get("sdk_function"), "sdk_function", path)
python_scope = _require_string_list(data.get("python_scope"), "python_scope", path)
rust_scope = _require_string_list(data.get("rust_scope"), "rust_scope", path)
entries_data = data.get("entries")
if not isinstance(entries_data, list):
raise ValueError(f"{path}: entries must be a list")
entries = tuple(
_load_entry(entry, index, path) for index, entry in enumerate(entries_data)
)
rust_only_data = data.get("rust_only_tests")
if not isinstance(rust_only_data, list):
raise ValueError(f"{path}: rust_only_tests must be a list")
rust_only_tests = tuple(
_load_rust_only_entry(entry, index, path) for index, entry in enumerate(rust_only_data)
)
return TestLedger(
sdk_function=sdk_function,
python_scope=python_scope,
rust_scope=rust_scope,
entries=entries,
rust_only_tests=rust_only_tests,
)

0
tests/rust-python-harness/strategies/__init__.py Normal file
View file

0
tests/rust-python-harness/strategies/unit_tests/__init__.py Normal file
View file

1084
tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json Normal file
View file

File diff suppressed because it is too large Load diff

101
tests/rust-python-harness/strategies/unit_tests/mapping_validator.py Normal file
View file

@ -0,0 +1,101 @@
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from ...shared.parity.ledger import TestLedger, load_ledger
from .python_runner import enumerate_python_tests
from .rust_runner import enumerate_rust_tests
REPO_ROOT = Path(__file__).resolve().parents[4]
LEDGER_ROOT = Path(__file__).parent / "ledgers"
def ledger_path_for(sdk_function: str) -> Path:
return LEDGER_ROOT / sdk_function / f"{sdk_function}_test_ledger.json"
@dataclass(frozen=True, slots=True)
class AuditReport:
missing_python_tests: tuple[str, ...]
stale_python_tests: tuple[str, ...]
missing_rust_tests: tuple[str, ...]
stale_rust_tests: tuple[str, ...]
@property
def is_clean(self) -> bool:
return not (
self.missing_python_tests
or self.stale_python_tests
or self.missing_rust_tests
or self.stale_rust_tests
)
def _ledger_python_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]:
grouping: dict[str, set[str]] = {path: set() for path in ledger.python_scope}
for entry in ledger.entries:
grouping.setdefault(entry.python_file, set()).add(entry.python_test)
return grouping
def _ledger_rust_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]:
grouping: dict[str, set[str]] = {path: set() for path in ledger.rust_scope}
for entry in ledger.entries:
if entry.status == "mapped":
grouping.setdefault(entry.rust_file, set()).add(entry.rust_test)
for rust_only in ledger.rust_only_tests:
grouping.setdefault(rust_only.rust_file, set()).add(rust_only.rust_test)
return grouping
def audit_ledger(ledger: TestLedger, repo_root: Path = REPO_ROOT) -> AuditReport:
missing_python: list[str] = []
stale_python: list[str] = []
for python_file, ledger_tests in _ledger_python_tests_by_file(ledger).items():
actual_tests = enumerate_python_tests(repo_root, python_file)
for missing in sorted(ledger_tests - actual_tests):
missing_python.append(f"{python_file}:{missing}")
for stale in sorted(actual_tests - ledger_tests):
stale_python.append(f"{python_file}:{stale}")
missing_rust: list[str] = []
stale_rust: list[str] = []
for rust_file, ledger_tests in _ledger_rust_tests_by_file(ledger).items():
actual_tests = enumerate_rust_tests(repo_root, rust_file)
for missing in sorted(ledger_tests - actual_tests):
missing_rust.append(f"{rust_file}:{missing}")
for stale in sorted(actual_tests - ledger_tests):
stale_rust.append(f"{rust_file}:{stale}")
return AuditReport(
missing_python_tests=tuple(missing_python),
stale_python_tests=tuple(stale_python),
missing_rust_tests=tuple(missing_rust),
stale_rust_tests=tuple(stale_rust),
)
@dataclass(frozen=True, slots=True)
class FunctionReport:
sdk_function: str
ledger: TestLedger | None
audit: AuditReport | None
@property
def has_ledger(self) -> bool:
return self.ledger is not None
@property
def is_clean(self) -> bool:
return self.audit is None or self.audit.is_clean
def build_function_report(sdk_function: str, repo_root: Path = REPO_ROOT) -> FunctionReport:
path = ledger_path_for(sdk_function)
if not path.exists():
return FunctionReport(sdk_function=sdk_function, ledger=None, audit=None)
ledger = load_ledger(path)
return FunctionReport(
sdk_function=sdk_function, ledger=ledger, audit=audit_ledger(ledger, repo_root)
)

22
tests/rust-python-harness/strategies/unit_tests/python_runner.py Normal file
View file

@ -0,0 +1,22 @@
from __future__ import annotations
import ast
from pathlib import Path
def enumerate_python_tests(repo_root: Path, relative_path: str) -> frozenset[str]:
source = (repo_root / relative_path).read_text(encoding="utf-8")
tree = ast.parse(source, filename=relative_path)
module_level: list[str] = []
for node in ast.iter_child_nodes(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_"):
module_level.append(node.name)
elif isinstance(node, ast.ClassDef):
for child in ast.iter_child_nodes(node):
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name.startswith(
"test_"
):
module_level.append(f"{node.name}::{child.name}")
return frozenset(module_level)

13
tests/rust-python-harness/strategies/unit_tests/rust_runner.py Normal file
View file

@ -0,0 +1,13 @@
from __future__ import annotations
import re
from pathlib import Path
_RUST_TEST_PATTERN = re.compile(
r"#\[(?:test|tokio::test)\][^\n]*\n(?:[^\n]*\n)*?\s*(?:async\s+)?fn\s+(\w+)\s*\("
)
def enumerate_rust_tests(repo_root: Path, relative_path: str) -> frozenset[str]:
source = (repo_root / relative_path).read_text(encoding="utf-8")
return frozenset(match.group(1) for match in _RUST_TEST_PATTERN.finditer(source))

2
tests/rust-python-harness/ui.py
View file

@ -119,7 +119,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]):
table = Table(box=box.ROUNDED, expand=True, title="Strategy × SDK function")
table.add_column("Strategy", ratio=3)
for label in ("ocr/aocr", "messages", "responses", "count_tokens"):
for label in SDK_FUNCTIONS:
table.add_column(label, justify="center", ratio=1)
for strategy in self.strategies:
cells = []

4
tests/rust-python-harness/unit_tests_rust/strategy.json
View file

@ -7,6 +7,8 @@
"ocr": {"coverage": "planned", "selectors": []},
"messages": {"coverage": "planned", "selectors": []},
"responses": {"coverage": "planned", "selectors": []},
"count_tokens": {"coverage": "planned", "selectors": []}
"count_tokens": {"coverage": "planned", "selectors": []},
"chat_completions": {"coverage": "planned", "selectors": []},
"transcription": {"coverage": "planned", "selectors": []}
}
}

4
tests/rust-python-harness/validate_sub_methods/strategy.json
View file

@ -7,6 +7,8 @@
"ocr": {"coverage": "planned", "selectors": []},
"messages": {"coverage": "planned", "selectors": []},
"responses": {"coverage": "planned", "selectors": []},
"count_tokens": {"coverage": "planned", "selectors": []}
"count_tokens": {"coverage": "planned", "selectors": []},
"chat_completions": {"coverage": "planned", "selectors": []},
"transcription": {"coverage": "planned", "selectors": []}
}
}

241
tests/test_litellm/integrations/azure_storage/test_azure_storage.py
View file

@ -1,10 +1,15 @@
import asyncio
import sys
import threading
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.integrations.azure_storage.azure_storage import (
AzureBlobStorageLogger,
_cached_credential_chain_token_provider,
)
from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType
from litellm.types.utils import StandardLoggingPayload
@ -25,6 +30,26 @@ def mock_gov_env_vars(mock_env_vars, monkeypatch):
monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", "core.usgovcloudapi.net")
@pytest.fixture
def workload_identity_env_vars(monkeypatch):
monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account")
monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container")
for unset in (
"AZURE_STORAGE_TENANT_ID",
"AZURE_STORAGE_CLIENT_ID",
"AZURE_STORAGE_CLIENT_SECRET",
"AZURE_STORAGE_ACCOUNT_KEY",
"AZURE_STORAGE_ENDPOINT_SUFFIX",
"AZURE_CLIENT_SECRET",
"AZURE_CREDENTIAL",
"AZURE_SCOPE",
):
monkeypatch.delenv(unset, raising=False)
monkeypatch.setenv("AZURE_CLIENT_ID", "workload-identity-client-id")
monkeypatch.setenv("AZURE_TENANT_ID", "workload-identity-tenant-id")
monkeypatch.setenv("AZURE_FEDERATED_TOKEN_FILE", "/var/run/secrets/azure/tokens/azure-identity-token")
@pytest.mark.asyncio
async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
"""
@ -32,17 +57,12 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
a payload to Azure Blob Storage using the 3-step process (create, append, flush).
"""
with (
patch(
"litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
) as mock_get_client,
patch(
"litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id"
) as mock_get_token,
patch("litellm.integrations.azure_storage.azure_storage.get_async_httpx_client") as mock_get_client,
patch("litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_from_entra_id") as mock_get_token,
):
# Create mock HTTP client
mock_http_client = AsyncMock()
mock_response = AsyncMock()
mock_response.raise_for_status = AsyncMock()
mock_response = MagicMock()
mock_http_client.put.return_value = mock_response
mock_http_client.patch.return_value = mock_response
mock_get_client.return_value = mock_http_client
@ -79,9 +99,7 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
put_call_args = mock_http_client.put.call_args
assert put_call_args[0][0] == f"{expected_base_url}?resource=file"
assert put_call_args[1]["headers"]["x-ms-version"] is not None
assert (
put_call_args[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
)
assert put_call_args[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
# Step 2: Append data
assert mock_http_client.patch.call_count == 2 # Called for append and flush
@ -89,9 +107,7 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
assert append_call[0][0] == f"{expected_base_url}?action=append&position=0"
assert append_call[1]["headers"]["x-ms-version"] is not None
assert append_call[1]["headers"]["Content-Type"] == "application/json"
assert (
append_call[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
)
assert append_call[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
assert "test-log-id-123" in append_call[1]["data"]
# Step 3: Flush data
@ -110,9 +126,7 @@ async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env
AZURE_STORAGE_ENDPOINT_SUFFIX must reach the Entra-ID REST upload path so a
sovereign-cloud account is addressed instead of the commercial dfs host.
"""
with patch(
"litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
) as mock_get_client:
with patch("litellm.integrations.azure_storage.azure_storage.get_async_httpx_client") as mock_get_client:
mock_http_client = AsyncMock()
mock_response = MagicMock()
mock_http_client.put.return_value = mock_response
@ -127,17 +141,10 @@ async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env
await logger.async_upload_payload_to_azure_blob_storage(test_payload)
expected_base_url = (
"https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json"
)
expected_base_url = "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json"
assert mock_http_client.put.call_args[0][0] == f"{expected_base_url}?resource=file"
assert (
mock_http_client.patch.call_args_list[0][0][0]
== f"{expected_base_url}?action=append&position=0"
)
assert mock_http_client.patch.call_args_list[1][0][0].startswith(
f"{expected_base_url}?action=flush"
)
assert mock_http_client.patch.call_args_list[0][0][0] == f"{expected_base_url}?action=append&position=0"
assert mock_http_client.patch.call_args_list[1][0][0].startswith(f"{expected_base_url}?action=flush")
@pytest.mark.asyncio
@ -148,9 +155,7 @@ async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars)
"""
fake_aio_module = MagicMock()
with patch.dict(
sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}
):
with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
logger = AzureBlobStorageLogger()
await logger.get_service_client()
@ -160,14 +165,180 @@ async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars)
)
@pytest.mark.asyncio
async def test_upload_authenticates_through_the_credential_chain_under_workload_identity(
workload_identity_env_vars,
):
build_provider = MagicMock(return_value=lambda: "workload-identity-token")
with patch( # test-quality-ok: REST client is created inside the method; assert emitted request headers
"litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
) as mock_get_client:
mock_http_client = AsyncMock()
mock_http_client.put.return_value = MagicMock()
mock_http_client.patch.return_value = MagicMock()
mock_get_client.return_value = mock_http_client
logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
await logger.async_upload_payload_to_azure_blob_storage({"id": "wif-log-id"})
build_provider.assert_called_once_with()
assert logger.azure_auth_token == "workload-identity-token"
sent_headers = [mock_http_client.put.call_args[1]["headers"]] + [
call[1]["headers"] for call in mock_http_client.patch.call_args_list
]
assert len(sent_headers) == 3
assert all(headers["Authorization"] == "Bearer workload-identity-token" for headers in sent_headers)
def test_default_chain_provider_is_storage_scoped_and_built_once_per_process():
_cached_credential_chain_token_provider.cache_clear()
with (
patch( # test-quality-ok: assert the default factory's fixed scope and credential type without constructing Azure SDK credentials
"litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_provider",
return_value=lambda: "chain-token",
) as mock_builder
):
first = _cached_credential_chain_token_provider()
second = _cached_credential_chain_token_provider()
_cached_credential_chain_token_provider.cache_clear()
assert first is second
assert first() == "chain-token"
mock_builder.assert_called_once_with(
azure_scope="https://storage.azure.com/.default",
azure_credential=AzureCredentialType.DefaultAzureCredential,
)
@pytest.mark.asyncio
async def test_chain_tokens_are_read_from_the_provider_on_every_refresh(
workload_identity_env_vars,
):
provider = MagicMock(side_effect=["chain-token-1", "chain-token-2"])
logger = AzureBlobStorageLogger(build_credential_chain_token_provider=MagicMock(return_value=provider))
await logger.set_valid_azure_ad_token()
first_token = logger.azure_auth_token
await logger.set_valid_azure_ad_token()
assert first_token == "chain-token-1"
assert logger.azure_auth_token == "chain-token-2"
assert provider.call_count == 2
@pytest.mark.asyncio
async def test_chain_token_read_yields_to_the_event_loop(workload_identity_env_vars):
"""
The chain walk is blocking I/O (IMDS probe, CLI subprocess), so reading the provider
inline would stall every request on the worker. Prove other coroutines run during the read.
"""
loop_was_free = threading.Event()
def provider() -> str:
if not loop_was_free.wait(timeout=5):
raise TimeoutError("the event loop never ran the observer while the token was being read")
return "chain-token"
async def observer():
loop_was_free.set()
logger = AzureBlobStorageLogger(build_credential_chain_token_provider=MagicMock(return_value=provider))
observer_task = asyncio.create_task(observer())
await logger.set_valid_azure_ad_token()
await observer_task
assert logger.azure_auth_token == "chain-token"
@pytest.mark.asyncio
async def test_empty_string_service_principal_vars_still_use_the_credential_chain(
workload_identity_env_vars, monkeypatch
):
for name in ("AZURE_STORAGE_TENANT_ID", "AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET"):
monkeypatch.setenv(name, "")
logger = AzureBlobStorageLogger(
build_credential_chain_token_provider=MagicMock(return_value=lambda: "workload-identity-token")
)
await logger.set_valid_azure_ad_token()
assert logger.azure_auth_token == "workload-identity-token"
@pytest.mark.asyncio
async def test_client_secret_auth_still_uses_the_storage_scoped_service_principal(mock_env_vars):
build_provider = MagicMock()
with (
patch( # test-quality-ok: assert the storage scope passed to the shared token factory without making an external auth call
"litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_from_entra_id",
return_value=lambda: "client-secret-token",
) as mock_entra_id
):
logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
await logger.set_valid_azure_ad_token()
assert logger.azure_auth_token == "client-secret-token"
build_provider.assert_not_called()
assert mock_entra_id.call_args.kwargs == {
"tenant_id": "test-tenant-id",
"client_id": "test-client-id",
"client_secret": "test-client-secret",
"scope": "https://storage.azure.com/.default",
}
@pytest.mark.parametrize(
"missing_var",
["AZURE_STORAGE_TENANT_ID", "AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET"],
)
@pytest.mark.asyncio
async def test_partially_configured_service_principal_still_names_the_missing_variable(
mock_env_vars, monkeypatch, missing_var
):
monkeypatch.delenv(missing_var)
build_provider = MagicMock()
logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
with pytest.raises(ValueError, match=f"Missing required environment variable: {missing_var}"):
await logger.set_valid_azure_ad_token()
build_provider.assert_not_called()
@pytest.mark.asyncio
async def test_account_key_auth_never_requests_a_token(workload_identity_env_vars, monkeypatch):
monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=")
file_client = MagicMock()
file_client.create_file = AsyncMock()
file_client.append_data = AsyncMock()
file_client.flush_data = AsyncMock()
directory_client = MagicMock()
directory_client.exists = AsyncMock(return_value=True)
directory_client.get_file_client = MagicMock(return_value=file_client)
file_system_client = MagicMock()
file_system_client.get_directory_client = MagicMock(return_value=directory_client)
service_client = MagicMock()
service_client.get_file_system_client = MagicMock(return_value=file_system_client)
fake_aio_module = MagicMock()
fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client)
build_provider = MagicMock()
with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
await logger.async_upload_payload_to_azure_blob_storage({"id": "account-key-log-id"})
build_provider.assert_not_called()
assert logger.azure_auth_token is None
file_client.flush_data.assert_awaited_once()
assert fake_aio_module.DataLakeServiceClient.call_args.kwargs["credential"] == "dGVzdC1rZXk="
@pytest.mark.asyncio
async def test_service_client_defaults_to_commercial_endpoint(mock_env_vars):
"""Unset AZURE_STORAGE_ENDPOINT_SUFFIX keeps the pre-existing commercial host"""
fake_aio_module = MagicMock()
with patch.dict(
sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}
):
with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
logger = AzureBlobStorageLogger()
await logger.get_service_client()

390
tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py
View file

@ -17,6 +17,7 @@ from unittest.mock import patch
import pytest
import litellm
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -55,15 +56,22 @@ def build_payload(
response_message: dict[str, Any] | None = None,
usage_object: dict[str, Any] | None = None,
model_parameters: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
model_group: str | None = None,
prompt_tokens: int = 4447,
) -> dict[str, Any]:
standard_logging_metadata: dict[str, Any] = {
**(metadata or {}),
**({"usage_object": usage_object} if usage_object is not None else {}),
}
return {
"standard_logging_object": {
"call_type": "acompletion",
"messages": [{"role": "user", "content": "hi"}] if messages is NOT_GIVEN else messages,
"response": {"choices": [{"message": response_message or {"role": "assistant", "content": "hello"}}]},
"model_parameters": model_parameters or {},
"metadata": {"usage_object": usage_object} if usage_object is not None else {},
"metadata": standard_logging_metadata,
"model_group": model_group,
"prompt_tokens": prompt_tokens,
"completion_tokens": 507,
"total_tokens": prompt_tokens + 507,
@ -244,6 +252,43 @@ def test_no_cache_keys_when_the_provider_reports_no_caching(logger: DataDogLLMOb
assert "non_cached_input_tokens" not in payload["metrics"]
def test_reasoning_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": 128}})
assert payload["metrics"]["reasoning_output_tokens"] == 128.0
def test_responses_reasoning_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, usage_object={"output_tokens_details": {"reasoning_tokens": 64}})
assert payload["metrics"]["reasoning_output_tokens"] == 64.0
def test_zero_reasoning_tokens_are_not_reported(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": 0}})
assert "reasoning_output_tokens" not in payload["metrics"]
def test_reasoning_tokens_come_from_the_spelling_that_reports_them(logger: DataDogLLMObsLogger) -> None:
"""A chat-details mapping without the count must not shadow the responses spelling that has it."""
payload = build(
logger,
usage_object={
"completion_tokens_details": {"accepted_prediction_tokens": 5},
"output_tokens_details": {"reasoning_tokens": 64},
},
)
assert payload["metrics"]["reasoning_output_tokens"] == 64.0
def test_boolean_reasoning_tokens_are_not_a_count(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": True}})
assert "reasoning_output_tokens" not in payload["metrics"]
def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, model_parameters={"tools": [TOOL_DEFINITION]})
@ -256,6 +301,340 @@ def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None:
]
def test_cost_tags_include_present_categories_and_dimensions(logger: DataDogLLMObsLogger) -> None:
payload = build(
logger,
metadata={
"user_api_key_user_id": "User 42",
"user_api_key_alias": "Primary Key",
"team_alias": "Platform",
"routing_decision": {
"tier": "premium",
"cause": "high_complexity",
"score": 0.91,
"escalated": True,
"signals": ["long prompt"],
"routed_model": "openai/gpt-5",
},
},
model_group="premium-models",
)
assert payload["tags"][-8:] == [
"team:platform",
"user:user_42",
"key_alias:primary_key",
"model_group:premium-models",
"router_tier:premium",
"router_cause:high_complexity",
"router_escalated:true",
"routed_model:openai/gpt-5",
]
assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == [
"team",
"user",
"key_alias",
"model_group",
"router_tier",
"router_cause",
"router_escalated",
"routed_model",
]
def test_missing_cost_tag_values_are_not_declared(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, metadata={"team_alias": "Platform"})
assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["team"]
assert not any(tag.startswith(("user:", "key_alias:", "model_group:")) for tag in payload["tags"])
def test_values_that_normalize_to_empty_are_not_tagged_or_declared(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, metadata={"user_api_key_user_id": "___", "user_api_key_alias": "!!!"}, model_group="tier-1")
assert not any(tag in ("user:", "key_alias:") for tag in payload["tags"])
assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["model_group"]
def test_a_valueless_tag_from_the_shared_builder_is_not_declared(logger: DataDogLLMObsLogger) -> None:
"""The team tag comes from the shared builder, which emits it bare when the alias normalizes away."""
payload = build(logger, metadata={"team_alias": "!!!"}, model_group="tier-1")
assert "team:" in payload["tags"]
assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["model_group"]
def test_router_fields_are_flattened(logger: DataDogLLMObsLogger) -> None:
payload = build(
logger,
metadata={
"routing_decision": {
"tier": "premium",
"cause": "high_complexity",
"score": 0.91,
"escalated": True,
"signals": ["secret prompt text"],
"routed_model": "openai/gpt-5",
}
},
model_group="premium-models",
)
assert payload["meta"]["metadata"]["router_tier"] == "premium"
assert payload["meta"]["metadata"]["router_cause"] == "high_complexity"
assert payload["meta"]["metadata"]["router_score"] == 0.91
assert payload["meta"]["metadata"]["router_escalated"] is True
assert payload["meta"]["metadata"]["router_signals"] == ["secret prompt text"]
assert payload["meta"]["metadata"]["routed_model"] == "openai/gpt-5"
def test_a_context_escalated_route_reports_as_escalated(logger: DataDogLLMObsLogger) -> None:
"""The router records a size-driven escalation under its own key, and it is still an escalation."""
payload = build(logger, metadata={"routing_decision": {"tier": "premium", "context_escalated": True}})
assert payload["meta"]["metadata"]["router_escalated"] is True
assert "router_escalated:true" in payload["tags"]
def test_a_routed_request_that_did_not_escalate_reports_false(logger: DataDogLLMObsLogger) -> None:
"""Without this the escalation dimension is absent on ordinary traffic, so nothing can group by it."""
payload = build(logger, metadata={"routing_decision": {"tier": "simple", "cause": "heuristic_scorer"}})
assert payload["meta"]["metadata"]["router_escalated"] is False
assert "router_escalated:false" in payload["tags"]
assert "router_escalated" in payload["meta"]["metadata"]["_dd"]["cost_tags"]
def test_a_request_that_never_reached_a_router_has_no_router_fields(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, model_group="premium-models")
assert "router_escalated" not in payload["meta"]["metadata"]
assert not any(tag.startswith("router_") for tag in payload["tags"])
def test_redacted_payload_keeps_metrics_and_removes_sensitive_fields(logger: DataDogLLMObsLogger) -> None:
payload = build_payload(
messages=[{"role": "user", "content": "secret prompt"}],
response_message={"role": "assistant", "content": "secret response"},
usage_object={"prompt_tokens_details": {"cached_tokens": 128}},
metadata={"routing_decision": {"tier": "premium", "signals": ["secret prompt text"]}},
model_parameters={"tools": [TOOL_DEFINITION]},
)
with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
with patch("asyncio.create_task"):
redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True)
redacted_payload = redacted_logger.redact_standard_logging_payload_from_model_call_details(payload)
result = json.loads(
safe_dumps(
redacted_logger.create_llm_obs_payload(
redacted_payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)
)
)
)
assert result["meta"]["input"]["messages"][0]["content"] == "redacted-by-litellm"
assert result["meta"]["output"]["messages"][0]["content"] == "redacted-by-litellm"
assert result["meta"]["metadata"]["router_tier"] == "premium"
assert "router_signals" not in result["meta"]["metadata"]
assert "routing_decision" not in result["meta"]["metadata"]
assert "tool_definitions" not in result["meta"]
assert result["metrics"]["cache_read_input_tokens"] == 128.0
assert result["metrics"]["total_cost"] == 0.02
def test_redaction_drops_the_routing_record_carried_in_metadata(logger: DataDogLLMObsLogger) -> None:
"""The whole routing record rides along in metadata, so dropping the flat copy alone leaks the prompt."""
with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
with patch("asyncio.create_task"):
redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True)
result = json.loads(
safe_dumps(
redacted_logger.create_llm_obs_payload(
build_payload(
metadata={
"routing_decision": {
"tier": "premium",
"cause": "keyword_rule",
"signals": ["secret prompt text"],
"matched_keyword": "secret keyword",
"escalation_keyword": "secret escalation",
}
}
),
datetime(2026, 9, 1, 12, 0, 0),
datetime(2026, 9, 1, 12, 0, 2),
)
)
)
assert "routing_decision" not in result["meta"]["metadata"]
assert result["meta"]["metadata"]["router_tier"] == "premium"
assert result["meta"]["metadata"]["router_cause"] == "keyword_rule"
assert "secret" not in safe_dumps(result["meta"]["metadata"])
def test_a_failure_span_redacts_its_messages(logger: DataDogLLMObsLogger) -> None:
"""The redaction hook only runs on success, so the failure span has to redact for itself."""
failed = build_payload(messages=[{"role": "user", "content": "secret prompt"}])
failed["standard_logging_object"]["status"] = "failure"
failed["standard_logging_object"]["response"] = None
failed["standard_logging_object"]["error_information"] = {"error_message": "boom", "error_class": "BadRequestError"}
with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
with patch("asyncio.create_task"):
redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True)
result = json.loads(
safe_dumps(
redacted_logger.create_llm_obs_payload(
failed, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)
)
)
)
assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
assert result["meta"]["output"]["messages"] == []
assert result["status"] == "error"
def test_excluding_messages_from_the_logging_payload_still_ships_the_span(logger: DataDogLLMObsLogger) -> None:
"""`standard_logging_payload_excluded_fields` deletes the key, and a span with no prompt is still a span."""
payload = build_payload()
del payload["standard_logging_object"]["messages"]
span = json.loads(
safe_dumps(
logger.create_llm_obs_payload(payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2))
)
)
assert span["meta"]["input"]["messages"] == []
assert span["metrics"]["total_cost"] == 0.02
def test_an_explicit_redaction_setting_survives_the_global_params(logger: DataDogLLMObsLogger) -> None:
"""Global params carry defaults for keys the operator never set, and those must not win."""
with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
with patch("asyncio.create_task"):
with patch.object( # test-quality-ok: the ctor reads this module global with no injection seam
litellm, "datadog_llm_observability_params", {}
):
configured_logger = DataDogLLMObsLogger(
turn_off_message_logging=True
) # test-quality-ok: verifies ctor setting
assert configured_logger.turn_off_message_logging is True
def _redacting_logger(
**kwargs: Any,
) -> DataDogLLMObsLogger: # test-quality-ok: shared test factory accepts init variants
with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
with patch("asyncio.create_task"):
return DataDogLLMObsLogger(**kwargs)
def _span_json(logger_under_test: DataDogLLMObsLogger, payload: dict[str, Any]) -> dict[str, Any]:
span = logger_under_test.create_llm_obs_payload(
payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)
)
return json.loads(safe_dumps(span))
def test_redaction_keeps_the_conversation_shape_without_its_content() -> None:
"""Roles and message count survive so the trace stays legible; contents and tool payloads do not."""
result = _span_json(
_redacting_logger(turn_off_message_logging=True),
build_payload(
messages=[
{"role": "user", "content": "secret prompt"},
{"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]},
],
response_message={"role": "assistant", "content": "secret response"},
),
)
assert result["meta"]["input"]["messages"] == [
{"role": "user", "content": "redacted-by-litellm"},
{"role": "assistant", "content": "redacted-by-litellm"},
]
assert result["meta"]["output"]["messages"] == [{"role": "assistant", "content": "redacted-by-litellm"}]
def test_redaction_drops_unrecognized_and_malformed_message_roles() -> None:
"""Caller-controlled role values must not bypass redaction or crash span creation."""
result = _span_json(
_redacting_logger(turn_off_message_logging=True),
build_payload(
messages=[
{"role": "SECRET-39402", "content": "hello"},
{"role": ["SECRET-39402"], "content": "hello"},
{"role": {"secret": "SECRET-39402"}, "content": "hello"},
{"role": "agent", "content": "hello"},
]
),
)
assert result["meta"]["input"]["messages"] == [
{"role": "", "content": "redacted-by-litellm"},
{"role": "", "content": "redacted-by-litellm"},
{"role": "", "content": "redacted-by-litellm"},
{"role": "agent", "content": "redacted-by-litellm"},
]
assert "SECRET-39402" not in safe_dumps(result)
def test_the_deprecated_message_logging_flag_engages_the_same_redaction() -> None:
"""The platform redacts for `message_logging is not True`, so this callback's own gate must agree."""
result = _span_json(
_redacting_logger(message_logging=False),
build_payload(
messages=[{"role": "user", "content": "secret prompt"}],
model_parameters={"tools": [TOOL_DEFINITION]},
metadata={"routing_decision": {"tier": "premium", "signals": ["secret prompt text"]}},
),
)
assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
assert "tool_definitions" not in result["meta"]
assert "routing_decision" not in result["meta"]["metadata"]
def test_a_truthy_redaction_setting_redacts_like_the_shared_hook() -> None:
"""The shared hook redacts on truthiness, so a config-provided string must not half-redact the span."""
result = _span_json(
_redacting_logger(turn_off_message_logging="yes"),
build_payload(messages=[{"role": "user", "content": "secret prompt"}]),
)
assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
def test_redaction_drops_every_prompt_carrying_metadata_record(logger: DataDogLLMObsLogger) -> None:
"""Tool arguments, retrieved text, and the guardrail's copy of the request ride in metadata records too."""
sensitive_metadata: dict[str, Any] = {
"requester_metadata": {"note": "secret prompt text"},
"prompt_management_metadata": {"prompt_id": "p1", "prompt_variables": {"topic": "secret"}},
"mcp_tool_call_metadata": {"name": "search", "arguments": {"query": "secret"}},
"vector_store_request_metadata": [{"query": "secret"}],
}
def sensitive_payload() -> dict[str, Any]:
payload = build_payload(metadata=sensitive_metadata)
payload["standard_logging_object"]["guardrail_information"] = [
{"guardrail_name": "g", "guardrail_request": {"messages": [{"content": "secret prompt"}]}}
]
return payload
redacted = _span_json(_redacting_logger(turn_off_message_logging=True), sensitive_payload())
unredacted = _span_json(logger, sensitive_payload())
assert "secret" not in safe_dumps(redacted["meta"]["metadata"])
for record in sensitive_metadata:
assert record not in redacted["meta"]["metadata"]
assert record in unredacted["meta"]["metadata"]
assert redacted["meta"]["metadata"]["guardrail_information"] is None
assert unredacted["meta"]["metadata"]["guardrail_information"] is not None
def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None:
"""The Anthropic surface declares tools unwrapped, with input_schema instead of parameters."""
payload = build(
@ -272,6 +651,15 @@ def test_meta_omits_tool_definitions_when_no_tools_were_offered(logger: DataDogL
assert "tool_definitions" not in build(logger)["meta"]
def test_a_ddtrace_integer_parent_id_is_forwarded_as_its_string(logger: DataDogLLMObsLogger) -> None:
"""ddtrace hands span ids as ints; dropping them detaches the span from its APM trace."""
kwargs = build_payload()
kwargs["litellm_params"]["metadata"]["parent_id"] = 8675309
start = datetime(2026, 9, 1, 12, 0, 0)
span = json.loads(safe_dumps(logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=2))))
assert span["parent_id"] == "8675309"
def test_unparseable_tool_arguments_are_preserved_rather_than_dropped(logger: DataDogLLMObsLogger) -> None:
"""A truncated argument string is still the only record of what the model tried to call."""
payload = build(

35
tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py Normal file
View file

@ -0,0 +1,35 @@
"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py."""
import json
from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
def _text_event(text: str) -> str:
return json.dumps(
{
"jsonrpc": "2.0",
"id": "req-1",
"result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]},
}
)
def _status_event() -> str:
return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}})
class TestA2AGuardrailHandlerStreamingScanKey:
def test_key_joins_the_text_of_every_message_event(self):
key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")])
assert key == StreamingScanKey(texts=("hello world",))
def test_events_without_text_leave_the_key_unchanged(self):
handler = A2AGuardrailHandler()
events = [_text_event("hello")]
assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events)
def test_unparseable_items_are_ignored(self):
key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"])
assert key.texts == ("hi",)

54
tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -13,6 +13,7 @@ import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
AnthropicMessagesHandler,
)
@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults:
}
later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)]
assert {"type": "text", "text": "Now fetch the page."} in later_blocks
class TestAnthropicMessagesHandlerStreamingScanKey:
"""get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
@staticmethod
def _sse(event_type, data):
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
def _text_delta(self, text):
return self._sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
)
def test_key_is_empty_before_any_text_arrives(self):
head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}})
key = AnthropicMessagesHandler().get_streaming_scan_key([head])
assert key == StreamingScanKey(texts=("",))
def test_key_accumulates_text_deltas(self):
key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")])
assert key.texts == ("hello world",)
assert key.stream_ended is False
def _stop(self, stop_reason):
return self._sse(
"message_delta",
{"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}},
)
def test_stop_without_tool_use_scans_the_same_payload(self):
handler = AnthropicMessagesHandler()
open_key = handler.get_streaming_scan_key([self._text_delta("hi")])
ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")])
assert ended_key.stream_ended is True
assert ended_key == open_key
def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self):
handler = AnthropicMessagesHandler()
tool_use = self._sse(
"content_block_start",
{
"type": "content_block_start",
"index": 1,
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}},
},
)
open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use])
ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")])
assert open_key == StreamingScanKey(texts=("hi",))
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key

47
tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py
View file

@ -27,6 +27,20 @@ def mock_gov_env_vars(mock_env_vars, monkeypatch):
monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", GOV_SUFFIX)
@pytest.fixture
def credential_chain_env_vars(monkeypatch):
monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account")
monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container")
for name in (
"AZURE_STORAGE_TENANT_ID",
"AZURE_STORAGE_CLIENT_ID",
"AZURE_STORAGE_CLIENT_SECRET",
"AZURE_STORAGE_ACCOUNT_KEY",
"AZURE_STORAGE_ENDPOINT_SUFFIX",
):
monkeypatch.delenv(name, raising=False)
def _make_backend() -> AzureBlobStorageBackend:
backend = AzureBlobStorageBackend()
backend.azure_auth_token = "mock-azure-ad-token"
@ -42,6 +56,29 @@ def _mock_upload_client() -> AsyncMock:
return client
@pytest.mark.asyncio
async def test_upload_file_with_credential_chain(credential_chain_env_vars):
client = _mock_upload_client()
build_provider = MagicMock(return_value=lambda: "workload-identity-token")
with patch( # test-quality-ok: the backend creates its REST client internally; assert the emitted authorization header
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=client
):
backend = AzureBlobStorageBackend(build_credential_chain_token_provider=build_provider)
storage_url = await backend.upload_file(
file_content=b"hello",
filename="report.json",
content_type="application/json",
path_prefix="logs",
file_naming_strategy="original_filename",
)
build_provider.assert_called_once_with()
assert storage_url == "https://test-account.blob.core.windows.net/test-container/logs/report.json"
assert client.put.call_args[1]["headers"]["Authorization"] == "Bearer workload-identity-token"
assert client.patch.call_count == 2
@pytest.mark.parametrize(
"env_fixture, expected_suffix",
[("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)],
@ -125,10 +162,7 @@ async def test_download_file_accepts_url_persisted_before_the_suffix_was_set(moc
)
assert content == b"file-bytes"
assert (
client.get.call_args[0][0]
== f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json"
)
assert client.get.call_args[0][0] == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json"
@pytest.mark.parametrize(
@ -178,10 +212,7 @@ async def test_download_file_drops_query_string_from_the_stored_url(mock_env_var
"https://test-account.blob.core.windows.net/test-container/logs/report.json?sig=redacted&se=2026"
)
assert (
client.get.call_args[0][0]
== "https://test-account.blob.core.windows.net/test-container/logs/report.json"
)
assert client.get.call_args[0][0] == "https://test-account.blob.core.windows.net/test-container/logs/report.json"
@pytest.mark.parametrize(

16
tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
View file

@ -8,6 +8,7 @@ import litellm
from litellm import get_model_info, supports_reasoning, supports_vision
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
from litellm.types.utils import (
ChatCompletionMessageToolCall,
@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id():
)
def test_get_fireworks_session_id_ignores_proxy_generated_session_id():
"""general_settings.missing_session_id: generate stamps a fresh id per request; sending it
as x-session-affinity would pin every request to a different node."""
assert (
get_fireworks_session_id(
{
"litellm_session_id": "generated-1",
"litellm_trace_id": "generated-1",
"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True},
}
)
is None
)
def test_handle_message_content_with_tool_calls():
config = FireworksAIConfig()
message = Message(

72
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -12,6 +12,7 @@ import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
from litellm.llms.openai.chat.guardrail_translation.handler import (
OpenAIChatCompletionsHandler,
)
@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded:
)
]
assert handler._check_streaming_has_ended(chunks) is True
class TestStreamingScanKey:
"""get_streaming_scan_key identifies what a sampled round would scan so the
unified hook can skip rounds that would re-scan already-cleared text"""
@staticmethod
def _chunk(content, finish_reason=None, index=0):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
return ModelResponseStream(
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)]
)
def test_key_carries_accumulated_text_and_open_stream(self):
handler = OpenAIChatCompletionsHandler()
key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")])
assert key == StreamingScanKey(texts=("hello",))
def test_chunks_without_text_leave_the_key_unchanged(self):
handler = OpenAIChatCompletionsHandler()
before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")])
after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)])
assert after == before
def test_finish_chunk_without_tool_calls_scans_the_same_payload(self):
handler = OpenAIChatCompletionsHandler()
open_key = handler.get_streaming_scan_key([self._chunk("hi")])
ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")])
assert open_key.stream_ended is False
assert ended_key.stream_ended is True
assert ended_key == open_key
def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self):
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
ModelResponseStream,
StreamingChoices,
)
handler = OpenAIChatCompletionsHandler()
tool_call = ChatCompletionDeltaToolCall(
id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}')
)
tool_chunk = ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)]
)
open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk])
ended_key = handler.get_streaming_scan_key(
[self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")]
)
assert open_key == StreamingScanKey(texts=("hi",))
assert ended_key.texts == ("hi",)
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key
def test_text_after_the_first_choice_finishes_still_changes_the_key(self):
handler = OpenAIChatCompletionsHandler()
first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)]
key_at_first_finish = handler.get_streaming_scan_key(first_done)
key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)])
assert key_at_first_finish.stream_ended is True
assert key_after_more_text.stream_ended is True
assert key_after_more_text != key_at_first_finish
def test_non_stream_items_are_ignored(self):
handler = OpenAIChatCompletionsHandler()
key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"])
assert key.texts == ("hi",)

90
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks:
dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"]
assert len(dones) == 1
assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy."
class TestOpenAIResponsesHandlerStreamingScanKey:
"""get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
@staticmethod
def _delta(sequence_number, text):
return {
"type": "response.output_text.delta",
"sequence_number": sequence_number,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": text,
}
def test_no_events_yields_no_key(self):
assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None
def test_key_accumulates_deltas_while_the_stream_is_open(self):
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")])
assert key == StreamingScanKey(texts=("hello",))
def test_typed_delta_events_accumulate_like_dicts(self):
from litellm.types.llms.openai import OutputTextDeltaEvent
events = [
OutputTextDeltaEvent(
type="response.output_text.delta",
item_id="msg_1",
output_index=0,
content_index=0,
delta=text,
sequence_number=i,
)
for i, text in enumerate(("hel", "lo"))
]
key = OpenAIResponsesHandler().get_streaming_scan_key(events)
assert key.texts == ("hello",)
assert key.stream_ended is False
def test_events_without_text_leave_the_key_unchanged(self):
handler = OpenAIResponsesHandler()
events = [self._delta(0, "hi")]
quiet = events + [{"type": "response.in_progress", "sequence_number": 1}]
assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events)
@staticmethod
def _completed(sequence_number, output):
return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}}
def test_completed_event_keys_on_the_final_output_text(self):
handler = OpenAIResponsesHandler()
message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
open_key = handler.get_streaming_scan_key([self._delta(0, "hi")])
ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])])
assert ended_key.stream_ended is True
assert ended_key == open_key
def test_completed_event_with_a_function_call_changes_the_key(self):
handler = OpenAIResponsesHandler()
message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}
open_key = handler.get_streaming_scan_key([self._delta(0, "hi")])
ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])])
assert ended_key.texts == ("hi",)
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key
def test_completed_event_reads_every_output_text_part(self):
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
item = GenericResponseOutputItem(
type="message",
id="msg_1",
status="completed",
role="assistant",
content=[
OutputText(type="output_text", text="one", annotations=[]),
OutputText(type="output_text", text="two", annotations=[]),
],
)
key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])])
assert key.texts == ("one", "two")
def test_output_item_done_round_is_never_deduped(self):
done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None

6
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -914,6 +914,7 @@ async def test_mcp_get_prompt_success():
)
mock_manager.get_prompt_from_server.assert_awaited_once_with(
server=server,
user_api_key_auth=user_api_key_auth,
prompt_name="hello",
arguments={"foo": "bar"},
mcp_auth_header={"Authorization": "token"},
@ -976,6 +977,7 @@ async def test_mcp_read_resource_success():
)
mock_manager.read_resource_from_server.assert_awaited_once_with(
server=server,
user_api_key_auth=user_api_key_auth,
url="https://example.com/resource",
mcp_auth_header={"Authorization": "token"},
extra_headers={"X-Test": "1"},
@ -8268,7 +8270,9 @@ class TestOboPreflightScopedToAllowedServers:
_, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key)
preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key)
preflight.assert_awaited_once_with(
server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None
)
@pytest.mark.asyncio

311
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -30,6 +30,7 @@ from mcp.types import (
TextResourceContents,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm.constants import MCP_METADATA_TIMEOUT
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@ -2270,7 +2271,9 @@ class TestMCPServerManager:
"""prompts/list on an OBO server must exchange the caller's bearer, not connect with none."""
server = self._token_exchange_server("te-prompts")
st = await self._capture_subject_token(
lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_prompts_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st == "subj-jwt"
@ -2279,7 +2282,9 @@ class TestMCPServerManager:
"""resources/list on an OBO server must exchange the caller's bearer."""
server = self._token_exchange_server("te-resources")
st = await self._capture_subject_token(
lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_resources_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st == "subj-jwt"
@ -2290,6 +2295,7 @@ class TestMCPServerManager:
st = await self._capture_subject_token(
lambda m: m.read_resource_from_server(
server=server,
user_api_key_auth=None,
url="https://up.example.com/r",
raw_headers={"authorization": "Bearer subj-jwt"},
)
@ -2307,7 +2313,9 @@ class TestMCPServerManager:
auth_type=MCPAuth.none,
)
st = await self._capture_subject_token(
lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_prompts_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st is None
@ -3254,7 +3262,7 @@ class TestMCPServerManager:
new_callable=AsyncMock,
return_value=mock_client,
):
prompts = await manager.get_prompts_from_server(server, add_prefix=True)
prompts = await manager.get_prompts_from_server(server, user_api_key_auth=None, add_prefix=True)
mock_client.list_prompts.assert_awaited_once()
assert len(prompts) == 1
@ -3289,6 +3297,7 @@ class TestMCPServerManager:
):
result = await manager.get_prompt_from_server(
server=server,
user_api_key_auth=None,
prompt_name="hello",
arguments={"tone": "casual"},
)
@ -3334,6 +3343,7 @@ class TestMCPServerManager:
):
result = await manager.get_resources_from_server(
server=server,
user_api_key_auth=None,
mcp_auth_header="auth",
extra_headers={"X-Test": "1"},
add_prefix=True,
@ -3391,6 +3401,7 @@ class TestMCPServerManager:
):
result = await manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=None,
mcp_auth_header="auth",
extra_headers=None,
add_prefix=False,
@ -3441,6 +3452,7 @@ class TestMCPServerManager:
) as mock_create_client:
result = await manager.read_resource_from_server(
server=server,
user_api_key_auth=None,
url="https://example.com/resource",
mcp_auth_header="auth",
extra_headers={"X-Test": "1"},
@ -11006,3 +11018,294 @@ class TestOpenApiHandlerRelaysUpstreamAuth:
assert result.isError is True
assert "upstream returned HTTP 503" in result.content[0].text
class TestLitellmAdmissionKeyIsNeverTheSubjectToken:
"""The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the
RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it
present, ``Authorization`` is the caller's own identity token and is exchanged as before."""
_ADMISSION_KEY: Final = "sk-litellm-virtual-key"
_USER_TOKEN: Final = "user-idp-jwt"
@staticmethod
def _token_exchange_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-server",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
)
@staticmethod
def _id_jag_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-server",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_id_jag,
client_id="cid",
client_secret="csec",
token_exchange_endpoint="https://idp.example.com/token",
id_jag_resource_token_endpoint="https://resource-as.example.com/token",
)
@staticmethod
def _recording_provider() -> MagicMock:
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
provider: Final = MagicMock()
provider.resolve_credentials = AsyncMock(
return_value=Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))
)
return provider
@staticmethod
def _subjects_seen_by(provider: MagicMock) -> list[str | None]:
return [
call.args[0].inbound_token.get_secret_value() if call.args[0].inbound_token else None
for call in provider.resolve_credentials.call_args_list
]
@staticmethod
def _manager_with_recording_client() -> MCPServerManager:
manager: Final = MCPServerManager()
client: Final = AsyncMock()
client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
client.list_prompts = AsyncMock(return_value=[])
client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[]))
manager._create_mcp_client = AsyncMock(return_value=client)
return manager
@staticmethod
def _subject_token_given_to_client(manager: MCPServerManager) -> str | None:
return manager._create_mcp_client.call_args.kwargs["subject_token"]
async def _call_tool_subject(self, server: MCPServer, oauth2_headers, raw_headers, user_api_key_auth):
manager: Final = self._manager_with_recording_client()
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=None,
user_api_key_auth=user_api_key_auth,
)
return self._subject_token_given_to_client(manager)
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag])
async def test_tools_call_with_only_the_litellm_key_has_no_subject(self, auth_type):
server = (
self._token_exchange_server("te-call")
if auth_type == MCPAuth.oauth2_token_exchange
else self._id_jag_server("jag-call")
)
subject_token = await self._call_tool_subject(
server,
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_rest_tools_call_with_only_the_litellm_key_has_no_subject(self):
"""The REST facade passes no oauth2_headers; the bearer is reached through raw_headers only."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-rest"),
oauth2_headers=None,
raw_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-split"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={
"X-LiteLLM-API-Key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token == self._USER_TOKEN
@pytest.mark.asyncio
async def test_tools_call_with_an_empty_x_litellm_api_key_has_no_subject(self):
"""Admission ignores an empty ``x-litellm-api-key`` and validates ``Authorization`` instead."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-empty-header"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"x-litellm-api-key": "", "authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_with_the_same_litellm_key_in_both_headers_has_no_subject(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-same-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={
"x-litellm-api-key": self._ADMISSION_KEY,
"authorization": f"Bearer {self._ADMISSION_KEY}",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_with_a_different_litellm_key_in_authorization_has_no_subject(self):
"""A second ``sk-`` virtual key next to ``x-litellm-api-key`` is still a gateway credential."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-second-key"),
oauth2_headers={"Authorization": "Bearer sk-another-virtual-key"},
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": "Bearer sk-another-virtual-key",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_exchanges_the_bearer_when_jwt_admission_left_api_key_unset(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-jwt"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={"authorization": f"Bearer {self._USER_TOKEN}"},
user_api_key_auth=UserAPIKeyAuth(api_key=None, user_id="alice"),
)
assert subject_token == self._USER_TOKEN
@pytest.mark.asyncio
async def test_tools_list_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
manager._fetch_tools_with_timeout = AsyncMock(return_value=[])
await manager._get_tools_from_server(
server=self._token_exchange_server("te-list-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_prompts_list_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
await manager.get_prompts_from_server(
server=self._token_exchange_server("te-prompts-key"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_resource_read_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
await manager.read_resource_from_server(
server=self._token_exchange_server("te-read-key"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
url=AnyUrl("file:///notes.txt"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_resource_read_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
manager: Final = self._manager_with_recording_client()
await manager.read_resource_from_server(
server=self._token_exchange_server("te-read-split"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
url=AnyUrl("file:///notes.txt"),
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
)
assert self._subject_token_given_to_client(manager) == self._USER_TOKEN
@pytest.mark.asyncio
async def test_openapi_call_never_hands_the_litellm_key_to_the_exchanger(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
server = MCPServer(
server_id="te-openapi",
name="te_openapi",
server_name="te_openapi",
url=None,
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
spec_path="https://api.example.com/openapi.json",
)
user_auth = UserAPIKeyAuth(api_key="hashed-key", user_id="alice")
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
mcp_auth_header=None,
user_api_key_auth=user_auth,
forwarded_headers=None,
)
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
mcp_auth_header=None,
user_api_key_auth=user_auth,
forwarded_headers=None,
)
assert self._subjects_seen_by(provider) == [None, self._USER_TOKEN]
@pytest.mark.asyncio
async def test_preflight_challenges_instead_of_exchanging_the_litellm_key(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
with pytest.raises(HTTPException) as exc_info:
await manager.preflight_token_exchange(
server=self._token_exchange_server("te-preflight-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert exc_info.value.status_code == 401
headers = exc_info.value.headers or {}
assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "")
assert self._subjects_seen_by(provider) == []
@pytest.mark.asyncio
async def test_preflight_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
await manager.preflight_token_exchange(
server=self._token_exchange_server("te-preflight-split"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
)
assert self._subjects_seen_by(provider) == [self._USER_TOKEN]

19
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py
View file

@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None):
@pytest.mark.asyncio
async def test_openai_moderation_streaming_default_uses_sampled_cadence():
"""Default config samples every 5th streamed chunk and runs a final aggregate
pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream
calls, plus 1 final = 3 total.
pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream
round is skipped because chunk 10 already scanned the full text, for 2 total calls
"""
import litellm
@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence():
):
pass
assert patched_make_request.await_count == 3, (
f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), "
assert patched_make_request.await_count == 2, (
f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; "
f"the end-of-stream round is skipped because chunk 10 already scanned the full text), "
f"got {patched_make_request.await_count}"
)
@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder
@pytest.mark.asyncio
async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled():
"""With streaming_end_of_stream_only=False and streaming_sampling_rate=2,
moderation runs every 2nd chunk during the stream, plus once more at end.
moderation runs every 2nd chunk during the stream. The terminal chunk scan covers
the final aggregate, for 3 total calls
"""
import litellm
@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab
):
pass
# 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls),
# plus the final aggregate pass after the stream ends (1 call) = 4 total.
assert patched_make_request.await_count == 4, (
f"Expected 4 moderation calls (3 sampled + 1 final aggregate), "
assert patched_make_request.await_count == 3, (
f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped "
f"because chunk 6 already scanned the full text), "
f"got {patched_make_request.await_count}"
)

9
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py
View file

@ -1703,8 +1703,8 @@ async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts:
@pytest.mark.parametrize(
("configured", "expected_calls"),
[
({}, 3),
({"streaming_sampling_rate": 2}, 6),
({}, 2),
({"streaming_sampling_rate": 2}, 5),
({"streaming_end_of_stream_only": True}, 1),
({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1),
],
@ -1712,7 +1712,10 @@ async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts:
async def test_streaming_params_from_config_control_output_scan_cadence(
configured: dict[str, object], expected_calls: int
) -> None:
"""10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once."""
"""10 chunks: default samples at 5 and 10, rate 2 samples 5 times, end-of-stream scans once.
The final pass is skipped because chunk 10 already scanned the complete output.
"""
handler = _initialize_from_config(mode="post_call", **configured)
assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls

18
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
View file

@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
@pytest.mark.asyncio
async def test_streaming_default_uses_sampled_cadence(self):
"""Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3."""
"""Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10
cover the full text, so the end-of-stream round is skipped and there are 2 calls
"""
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
):
pass
assert mock_post.await_count == 3, (
f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), "
assert mock_post.await_count == 2, (
f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; "
f"the end-of-stream round is skipped because chunk 10 already scanned the full text), "
f"got {mock_post.await_count}"
)
for call in mock_post.await_args_list:
@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
@pytest.mark.asyncio
async def test_streaming_sampling_rate_override(self):
"""sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls."""
"""sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so
the end-of-stream round is skipped and there are 3 calls
"""
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
):
pass
assert mock_post.await_count == 4, (
f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), "
assert mock_post.await_count == 3, (
f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped "
f"because chunk 6 already scanned the full text), "
f"got {mock_post.await_count}"
)

1151
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

File diff suppressed because it is too large Load diff

268
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket:
assert recorded[0]["guardrail_name"] == "audit-recorder"
assert recorded[0]["guardrail_status"] == "success"
assert request_data["metadata"]["user_api_key_user_id"] == "user-1"
class _ScanCountingGuardrail(CustomGuardrail):
"""Pass-through guardrail that records every response-side scan payload."""
def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False):
super().__init__(guardrail_name="scan-counter")
self.streaming_sampling_rate = sampling_rate
self.streaming_end_of_stream_only = end_of_stream_only
self.streaming_buffer_until_moderated = buffer_until_moderated
self.guardrail_config = {}
self.scans: tuple[dict[str, object], ...] = ()
def should_run_guardrail(self, data, event_type): # type: ignore[override]
return True
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
self.scans = (
*self.scans,
{
"texts": list(inputs.get("texts") or []),
"tool_calls": list(inputs.get("tool_calls") or []),
"model": inputs.get("model"),
},
)
return inputs
def _responses_delta(sequence_number, text):
return {
"type": "response.output_text.delta",
"sequence_number": sequence_number,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": text,
}
def _responses_tail(sequence_number, text):
return [
{
"type": "response.output_text.done",
"sequence_number": sequence_number,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"text": text,
},
{
"type": "response.completed",
"sequence_number": sequence_number + 1,
"response": {
"model": "gpt-5.6",
"output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}],
},
},
]
class TestStreamingScanDedup:
"""A sampled round whose scan payload matches the previous round (or carries
no text yet) is skipped, so a stream is never re-scanned for output the
guardrail already cleared. Regression for LIT-6692."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self, monkeypatch):
monkeypatch.setattr(
unified_module,
"endpoint_guardrail_translation_mappings",
load_guardrail_translation_mappings(),
)
@pytest.mark.asyncio
async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 3
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
@pytest.mark.asyncio
async def test_chat_round_with_unchanged_text_is_skipped(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [
_stream_chunk("a"),
_stream_chunk("b"),
_stream_chunk("c"),
_stream_chunk(None),
_stream_chunk(None),
_stream_chunk(None),
_stream_chunk("d", finish_reason="stop"),
]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 7
assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]]
@pytest.mark.asyncio
async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 4
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
@pytest.mark.asyncio
async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self):
from litellm.types.utils import ChatCompletionDeltaToolCall, Function
guardrail = _ScanCountingGuardrail(sampling_rate=3)
tool_call = ChatCompletionDeltaToolCall(
id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}')
)
finish = ModelResponseStream(
choices=[
StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls")
]
)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 4
assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]]
assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"]
@pytest.mark.asyncio
async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [
_stream_chunk("a", index=0),
_stream_chunk("x", index=1),
_stream_chunk("b", finish_reason="stop", index=0),
_stream_chunk("y", index=1),
_stream_chunk("z", finish_reason="stop", index=1),
]
await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(guardrail.scans) == 2
assert any("yz" in text for text in guardrail.scans[-1]["texts"])
@pytest.mark.asyncio
async def test_responses_completed_event_on_sampled_index_is_scanned_once(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(8)]
full_text = "".join(f"t{i}" for i in range(8))
chunks = deltas + _responses_tail(8, full_text)
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 10
assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]]
assert guardrail.scans[-1]["model"] == "gpt-5.6"
@pytest.mark.asyncio
async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
chunks = deltas + _responses_tail(5, "t0t1t2t3t4")
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 7
assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]]
@pytest.mark.asyncio
async def test_responses_completed_carrying_a_function_call_is_still_scanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
completed = {
"type": "response.completed",
"sequence_number": 5,
"response": {
"model": "gpt-5.6",
"output": [
{"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]},
{
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "get_weather",
"arguments": '{"city": "Paris"}',
"status": "completed",
},
],
},
}
chunks = deltas + [completed]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 6
assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]]
assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"]
@pytest.mark.asyncio
async def test_responses_round_with_unchanged_text_is_skipped(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)]
chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4")
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 12
assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},)
@pytest.mark.asyncio
async def test_responses_tool_call_done_event_is_still_scanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=2)
tool_call_done = {
"type": "response.output_item.done",
"sequence_number": 1,
"output_index": 1,
"item": {
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "get_weather",
"arguments": '{"city": "Paris"}',
"status": "completed",
},
}
chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi")
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 4
assert len(guardrail.scans) == 2
assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"]
assert guardrail.scans[1]["texts"] == ["hi"]
@pytest.mark.asyncio
async def test_anthropic_skips_empty_round_and_terminal_duplicate(self):
guardrail = _ScanCountingGuardrail(sampling_rate=2)
chunks = _anthropic_message_chunks(["hello ", "world"])
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
assert out == chunks
assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]]
@pytest.mark.asyncio
async def test_end_of_stream_only_still_scans_exactly_once(self):
guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True)
chunks = _anthropic_message_chunks(["hello ", "world"])
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
assert out == chunks
assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]]
@pytest.mark.asyncio
async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self):
guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert out == chunks
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]

176
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import (
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
TRUSTED_CALLBACK_VARS_FIELD,
)
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
from litellm.types.utils import CredentialItem
@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge():
}
assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"]
def _request_for(path: str) -> MagicMock:
request = MagicMock(spec=Request)
request.scope = {"path": path}
request.url = MagicMock()
request.url.path = path
request.url.__str__.return_value = f"http://localhost{path}"
request.method = "POST"
request.query_params = {}
request.headers = {"Content-Type": "application/json"}
request.client = MagicMock()
request.client.host = "127.0.0.1"
return request
def _spend_log_session_id(data: dict[str, object]) -> str:
"""Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log
metadata = data["metadata"]
assert isinstance(metadata, dict)
litellm_params = get_litellm_params(
litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None,
litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None,
metadata=metadata,
)
trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id(
logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"),
litellm_params=litellm_params,
)
return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id})
@pytest.mark.asyncio
@pytest.mark.parametrize("request_correlation_in_logs", [False, True])
async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree(
monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool
):
"""Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs
must be the same generated id, so cross-referencing the two by session_id works. The id is marked
as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it."""
monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs)
data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
updated = await add_litellm_data_to_request(
data=data,
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "generate"},
)
callback_session_id = updated["metadata"]["session_id"]
assert isinstance(callback_session_id, str) and len(callback_session_id) == 36
assert _spend_log_session_id(updated) == callback_session_id
assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True
assert get_fireworks_session_id(
{"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]}
) is None
@pytest.mark.asyncio
async def test_missing_session_id_unset_keeps_legacy_divergence():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
)
assert "session_id" not in updated["metadata"]
assert "litellm_session_id" not in updated
assert _spend_log_session_id(updated) == "per-call-random-trace-id"
@pytest.mark.asyncio
async def test_missing_session_id_generate_reuses_traceparent_trace_id():
"""A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it."""
request = _request_for("/v1/chat/completions")
request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "generate"},
)
assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736"
@pytest.mark.asyncio
@pytest.mark.parametrize("policy", ["generate", "reject"])
async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str):
request = _request_for("/v1/chat/completions")
request.headers = {"x-litellm-session-id": "client-session-1"}
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": policy},
)
assert updated["litellm_session_id"] == "client-session-1"
assert updated["metadata"]["session_id"] == "client-session-1"
assert _spend_log_session_id(updated) == "client-session-1"
assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"]
assert (
get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]})
== "client-session-1"
)
@pytest.mark.asyncio
async def test_missing_session_id_reject_accepts_body_metadata_session_id():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "reject"},
)
assert updated["metadata"]["session_id"] == "body-session-1"
@pytest.mark.asyncio
async def test_missing_session_id_reject_returns_400_without_session_id():
with pytest.raises(ProxyException) as exc_info:
await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "reject"},
)
assert exc_info.value.code == "400"
assert exc_info.value.param == "session_id"
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"])
async def test_missing_session_id_policy_skips_non_inference_routes(path: str):
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o"},
request=_request_for(path),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "reject"},
)
assert "session_id" not in updated["metadata"]
@pytest.mark.asyncio
async def test_missing_session_id_unknown_value_is_ignored():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "typo"},
)
assert "session_id" not in updated["metadata"]

90
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -16,7 +16,7 @@ import litellm
from litellm import Router
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.router_strategy.complexity_router.complexity_router import (
_CLASSIFICATION_CURRENT_MESSAGE_ONLY,
_CLASSIFICATION_WITH_CONVERSATION,
@ -4274,6 +4274,26 @@ class TestSessionAffinity:
assert first.model == "o1-preview"
assert second.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config):
"""A session id the proxy generated for a request that had none is per request, so
it must not create a pin even with session_affinity enabled."""
mock_router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=session_affinity_config,
)
request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}}
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
)
second = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert first.model == "o1-preview"
assert second.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config):
"""Regression: session_affinity=True is the opt-in, so a shared session_id reuses the
@ -10777,6 +10797,9 @@ class TestModalityRouting:
("custom_tiers_walk", "premium-model", "modality_escalation"),
("pin_kept_bypasses", "text-cheap", "session_affinity_pin"),
("pin_replacement_gated", "vision-big", "modality_escalation"),
("pin_override_escalates", "vision-mid", "modality_pin_override"),
("pin_override_same_tier", "vision-cheap", "modality_pin_override"),
("pin_override_inert_without_modality_routing", "text-cheap", "session_affinity_pin"),
("adaptive_pick_rewritten", "vision-mid", "modality_escalation"),
],
)
@ -10827,7 +10850,7 @@ class TestModalityRouting:
messages = [
{"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]}
]
elif path in ("pin_kept_bypasses", "pin_replacement_gated"):
elif path.startswith(("pin_kept", "pin_replacement", "pin_override")):
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
mock_router_instance.cache = cache
@ -10839,6 +10862,13 @@ class TestModalityRouting:
messages = [
{"role": "user", "content": [{"type": "text", "text": "LITELLM ESCALATE describe this"}, IMG_PART]}
]
elif path == "pin_override_same_tier":
config["modality_pin_override"] = True
config["tiers"]["SIMPLE"] = ["text-cheap", "vision-cheap"]
vision["vision-cheap"] = True
elif path == "pin_override_inert_without_modality_routing":
config["modality_routing"] = False
config["modality_pin_override"] = path.startswith("pin_override")
elif path == "adaptive_pick_rewritten":
config["adaptive"] = True
mock_router_instance.model_list = []
@ -11005,4 +11035,60 @@ class TestModalityRouting:
from litellm.router_strategy.complexity_router.complexity_router import _decision_is_pinnable
assert _decision_is_pinnable({"cause": "modality_escalation"}) is False
assert _decision_is_pinnable({"cause": "modality_pin_override"}) is False
assert _decision_is_pinnable({"cause": "heuristic_scorer"}) is True
@pytest.mark.asyncio
async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance):
"""The override is for one request: the session keeps the model it was pinned to."""
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
mock_router_instance.cache = cache
router = self._router(
mock_router_instance,
{
"tiers": dict(self.BASE_TIERS),
"modality_routing": True,
"modality_pin_override": True,
"session_affinity": True,
},
dict(self.BASE_VISION),
)
request_kwargs = {"metadata": {"session_id": "s1"}}
image_turn = await router.async_pre_routing_hook(
model="m", request_kwargs=request_kwargs, messages=self.IMAGE_MESSAGE
)
assert image_turn.model == "vision-mid"
assert image_turn.routing_decision["cause"] == "modality_pin_override"
assert "modality_escalated_from:SIMPLE" in image_turn.routing_decision["signals"]
assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}
text_turn = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=[{"role": "user", "content": "hi"}]
)
assert text_turn.model == "text-cheap"
assert text_turn.routing_decision["cause"] == "session_affinity_pin"
@pytest.mark.asyncio
async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance):
"""The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin."""
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
mock_router_instance.cache = cache
router = self._router(
mock_router_instance,
{
"tiers": {"SIMPLE": "text-cheap", "COMPLEX": "text-big"},
"modality_routing": True,
"modality_pin_override": True,
"session_affinity": True,
},
{"text-cheap": False, "text-big": False},
)
with pytest.raises(litellm.BadRequestError, match="no model"):
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE
)
assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}

43
tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py
View file

@ -7,7 +7,7 @@ import json
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key():
assert filtered[0]["model_info"]["id"] == "deployment-2"
@pytest.mark.asyncio
async def test_proxy_generated_session_id_does_not_pin_a_deployment():
"""A session id the proxy generated for a request that had none is per request, so a
pin stored under it must be ignored and none must be written."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=123,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=True,
)
healthy_deployments = [
{"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}},
{"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}},
]
await cache.async_set_cache(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"),
{"model_id": "deployment-2"},
)
request_kwargs = {
"metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}
}
filtered = await callback.async_filter_deployments(
model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs
)
await callback.async_pre_call_deployment_hook(
kwargs={
"metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"},
"model_info": {"id": "deployment-1"},
},
call_type=None,
)
assert len(filtered) == 2
assert await cache.async_get_cache(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1")
) == {"model_id": "deployment-2"}
MOCK_RESPONSES_API_RESPONSE = {
"id": "resp_mock-resp-456",
"object": "response",

86
tests/test_litellm/test_check_py310_typing_imports.py Normal file
View file

@ -0,0 +1,86 @@
import sys
from pathlib import Path
from typing import Final
_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests"
sys.path.insert(0, str(_CODE_COVERAGE_DIR)) # test-quality-ok: required to import checker from its source directory
import check_py310_typing_imports as checker # noqa: E402 # load checker from its source directory
def _scan(tmp_path: Path, source: str) -> tuple[object, ...]:
file_path = tmp_path / "fixture.py"
file_path.write_text(source, encoding="utf-8")
return checker.scan_file(file_path)
def test_typing_import_flags_python_311_name(tmp_path: Path) -> None:
violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n")
assert tuple(violation.name for violation in violations) == ("NotRequired",)
def test_typing_extensions_import_passes(tmp_path: Path) -> None:
assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == ()
def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None:
violations = _scan(tmp_path, "import typing\nx: typing.Self\n")
assert tuple(violation.name for violation in violations) == ("Self",)
def test_version_guarded_typing_import_passes(tmp_path: Path) -> None:
source = (
"import sys\n"
"if sys.version_info >= (3, 11):\n"
" from typing import NotRequired\n"
"else:\n"
" from typing_extensions import NotRequired\n"
)
assert _scan(tmp_path, source) == ()
def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None:
source = (
"import sys\n"
"if sys.version_info >= (3, 11):\n"
" from typing_extensions import NotRequired\n"
"else:\n"
" from typing import NotRequired\n"
)
violations = _scan(tmp_path, source)
assert tuple(violation.name for violation in violations) == ("NotRequired",)
def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None:
source = (
"import sys\n"
"if sys.version_info < (3, 11):\n"
" from typing_extensions import NotRequired\n"
"else:\n"
" from typing import NotRequired\n"
)
assert _scan(tmp_path, source) == ()
def test_nearest_if_controls_version_guard(tmp_path: Path) -> None:
source = (
"if sys.version_info >= (3, 11):\n"
" from typing import Self\n"
" x = 1\n"
"if True:\n"
" from typing import Self\n"
)
violations = _scan(tmp_path, source)
assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),)
def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None:
file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py"
file_path.parent.mkdir(parents=True)
file_path.write_text("from typing import NotRequired\n", encoding="utf-8")
violations = checker.scan_directory(tmp_path)
assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),)
def test_python_310_typing_name_passes(tmp_path: Path) -> None:
assert _scan(tmp_path, "from typing import Optional\n") == ()

71
tests/test_rust_python_harness.py
View file

@ -8,14 +8,24 @@ import pytest
catalog = importlib.import_module("tests.rust-python-harness.catalog")
cli = importlib.import_module("tests.rust-python-harness.cli")
ledger_module = importlib.import_module("tests.rust-python-harness.shared.parity.ledger")
mapping_validator = importlib.import_module(
"tests.rust-python-harness.strategies.unit_tests.mapping_validator"
)
models = importlib.import_module("tests.rust-python-harness.models")
runner = importlib.import_module("tests.rust-python-harness.runner")
ui = importlib.import_module("tests.rust-python-harness.ui")
load_catalog = catalog.load_catalog
load_ledger = ledger_module.load_ledger
ledger_path_for = mapping_validator.ledger_path_for
REPO_ROOT = mapping_validator.REPO_ROOT
audit_ledger = mapping_validator.audit_ledger
build_function_report = mapping_validator.build_function_report
_pick_values = cli._pick_values
_coverage_pytest_args = cli._coverage_pytest_args
_select = cli._select
_validate_ledger = cli._validate_ledger
CaseResult = models.CaseResult
Coverage = models.Coverage
HarnessCase = models.HarnessCase
@ -56,13 +66,14 @@ def _manifest() -> dict[str, object]:
}
def test_should_load_the_three_harness_strategies_in_order() -> None:
def test_should_load_the_four_harness_strategies_in_order() -> None:
strategies = load_catalog()
assert [strategy.id for strategy in strategies] == [
"e2e_fuzz_tests",
"unit_tests_rust",
"validate_sub_methods",
"existing_e2e_test_sdk",
]
assert all(
tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS
@ -94,6 +105,8 @@ def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> Non
True,
),
("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False),
("tests/ocr_tests/", "tests/ocr_tests/test_ocr_mistral.py::test_one", True),
("tests/ocr_tests/", "tests/other_tests/test_ocr_mistral.py::test_one", False),
],
)
def test_should_match_pytest_file_and_node_selectors(
@ -113,6 +126,13 @@ def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None:
assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",)
def test_should_treat_an_existing_folder_selector_as_runnable(tmp_path: Path) -> None:
(tmp_path / "tests" / "ocr_tests").mkdir(parents=True)
case = _case(selectors=("tests/ocr_tests/",))
assert runnable_selectors((case,), tmp_path) == ("tests/ocr_tests/",)
def test_should_mark_planned_and_not_applicable_cases_without_running() -> None:
planned = CaseResult(case=_case(coverage=Coverage.PLANNED))
not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE))
@ -229,8 +249,53 @@ def test_should_report_confidence_for_each_sdk_section() -> None:
}
assert scores["responses"].verified_strategies == 1
assert scores["responses"].required_strategies == 3
assert scores["responses"].percentage == 33
assert scores["responses"].required_strategies == 4
assert scores["responses"].percentage == 25
assert scores["responses"].level.value == "MEDIUM"
assert scores["count_tokens"].percentage == 0
assert scores["count_tokens"].level.value == "LOW"
def test_should_report_no_ledger_for_a_function_without_one() -> None:
report = build_function_report("messages", repo_root=REPO_ROOT)
assert report.has_ledger is False
assert report.is_clean is True
def test_should_report_ocr_ledger_stats_and_a_clean_audit() -> None:
ledger = load_ledger(ledger_path_for("ocr"))
report = build_function_report("ocr", repo_root=REPO_ROOT)
assert report.has_ledger is True
assert report.ledger.mapped_count == ledger.mapped_count
assert report.ledger.total_count == ledger.total_count
assert report.is_clean is True
def test_should_scope_validate_ledger_to_the_requested_function(
capsys: pytest.CaptureFixture[str],
) -> None:
exit_code = _validate_ledger({"messages"})
captured = capsys.readouterr()
assert exit_code == 0
assert "messages" in captured.out
assert "no ledger yet" in captured.out
assert "ocr" not in captured.out
def test_should_have_every_python_and_rust_ocr_test_accounted_for_in_the_ledger() -> None:
ledger = load_ledger(ledger_path_for("ocr"))
report = audit_ledger(ledger, repo_root=REPO_ROOT)
assert report.is_clean, (
"\nOCR test-parity ledger is out of sync with the live test files.\n"
f"Ledger references a Python test that no longer exists: {list(report.missing_python_tests)}\n"
f"Python test exists but is not tracked in the ledger: {list(report.stale_python_tests)}\n"
f"Ledger references a Rust test that no longer exists: {list(report.missing_rust_tests)}\n"
f"Rust test exists but is not tracked in the ledger: {list(report.stale_rust_tests)}\n"
)

4
type-discipline-budget.json
View file

@ -1,6 +1,6 @@
{
"LIT001": {
"limit": 22334
"limit": 22330
},
"LIT002": {
"limit": 26763
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16480
"limit": 16478
},
"LIT011": {
"limit": 5520

38
ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
View file

@ -880,6 +880,44 @@ describe("ComplexityRouterConfig modality panel", () => {
expect(screen.getByRole("switch", { name: "Route image requests to vision-capable models" })).toBeChecked();
});
// The backend ignores modality_pin_override unless modality_routing is on, so offering it while
// image routing is off would let an operator save a flag that does nothing.
it("disables the pin-override switch while image routing is off", () => {
const onChange = vi.fn();
renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Modality Routing"));
const override = screen.getByRole("switch", { name: "Override session pin for image requests" });
expect(override).toHaveAttribute("aria-disabled", "true");
fireEvent.click(override);
expect(onChange).not.toHaveBeenCalled();
});
it("writes modality_pin_override through onChange once image routing is on", () => {
const onChange = vi.fn();
const value = { ...defaultValue, modality_routing: true };
renderWithProviders(<ComplexityRouterConfig {...baseProps} value={value} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Modality Routing"));
const override = screen.getByRole("switch", { name: "Override session pin for image requests" });
expect(override).not.toBeChecked();
fireEvent.click(override);
expect(onChange).toHaveBeenCalledWith({ ...value, modality_pin_override: true });
});
it("renders a stored modality_pin_override=true as on", () => {
renderWithProviders(
<ComplexityRouterConfig
{...baseProps}
value={{ ...defaultValue, modality_routing: true, modality_pin_override: true }}
/>,
);
fireEvent.click(screen.getByText("Advanced: Modality Routing"));
expect(screen.getByRole("switch", { name: "Override session pin for image requests" })).toBeChecked();
});
});
describe("ComplexityRouterConfig affinity panel", () => {

1
ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
View file

@ -410,6 +410,7 @@ export interface ComplexityRouterConfigValue {
classification_mode?: ClassificationMode;
session_affinity?: boolean;
modality_routing?: boolean;
modality_pin_override?: boolean;
deployment_affinity?: boolean;
/** Plan-mode floor as a tier ROW ID, unset meaning off. The wire carries the row's name. */
plan_mode_min_tier?: string;

50
ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx
View file

@ -7,20 +7,36 @@ import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
export const ModalityRoutingControls: React.FC<{
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
}> = ({ value, onChange }) => (
<>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={value.modality_routing ?? false}
onCheckedChange={(modalityRouting) => onChange({ ...value, modality_routing: modalityRouting })}
aria-label="Route image requests to vision-capable models"
/>
<strong className="font-semibold">Route image requests to vision-capable models</strong>
</div>
<span className="block text-xs text-muted-foreground">
Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default
model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are replaced,
and a kept session pin still wins.
</span>
</>
);
}> = ({ value, onChange }) => {
const modalityRouting = value.modality_routing ?? false;
return (
<>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={modalityRouting}
onCheckedChange={(nextModalityRouting) => onChange({ ...value, modality_routing: nextModalityRouting })}
aria-label="Route image requests to vision-capable models"
/>
<strong className="font-semibold">Route image requests to vision-capable models</strong>
</div>
<span className="block text-xs mb-3 text-muted-foreground">
Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default
model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are
replaced, and a kept session pin still wins unless you turn on the override below.
</span>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={value.modality_pin_override ?? false}
onCheckedChange={(modalityPinOverride) => onChange({ ...value, modality_pin_override: modalityPinOverride })}
disabled={!modalityRouting}
aria-label="Override session pin for image requests"
/>
<strong className="font-semibold">Override session pin for image requests</strong>
</div>
<span className="block text-xs text-muted-foreground">
Route an image turn to a capable model even when the session is pinned to one that cannot take images. The pin
is kept, so the next text turn goes back to it. Needs image routing turned on.
</span>
</>
);
};

38
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
View file

@ -590,6 +590,44 @@ describe("AddAutoRouterTab", () => {
});
});
it("writes both modality flags as false into the create payload when the panel stays untouched", async () => {
const user = userEvent.setup();
vi.mocked(getMissingTiersError).mockReturnValue(null);
renderWithProviders(<Harness />);
fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } });
await user.click(screen.getByRole("button", { name: /add auto router/i }));
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
modality_routing: false,
modality_pin_override: false,
});
});
it("carries the pin override through to the create payload once image routing unlocks it", async () => {
const user = userEvent.setup();
vi.mocked(getMissingTiersError).mockReturnValue(null);
renderWithProviders(<Harness />);
fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } });
expandDetailedConfiguration();
await user.click(screen.getByText("Advanced: Modality Routing"));
await user.click(await screen.findByRole("switch", { name: "Route image requests to vision-capable models" }));
await user.click(await screen.findByRole("switch", { name: "Override session pin for image requests" }));
await user.click(screen.getByRole("button", { name: /add auto router/i }));
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
modality_routing: true,
modality_pin_override: true,
});
});
// Custom is the escape hatch, not the headline choice, so it's listed after every bundled preset
// rather than first.
it("lists Custom Configuration after the bundled presets", () => {

1
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
View file

@ -358,6 +358,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
classifierFallback: complexityRouterConfig.classifier_fallback,
sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY,
modalityRouting: complexityRouterConfig.modality_routing ?? false,
modalityPinOverride: complexityRouterConfig.modality_pin_override ?? false,
deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
customTechnicalKeywords,
keywordTierRules,

9
ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
View file

@ -56,6 +56,7 @@ describe("buildComplexityRouterConfig", () => {
session_affinity: false,
deployment_affinity: true,
modality_routing: false,
modality_pin_override: false,
escalation_keywords: ["LITELLM ESCALATE"],
};
expect(config).toEqual(expected);
@ -270,6 +271,14 @@ describe("buildComplexityRouterConfig", () => {
expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: false }).modality_routing).toBe(false);
});
it("writes modality_pin_override explicitly both ways, so the stored config never relies on the backend default", () => {
expect(buildComplexityRouterConfig({ ...baseParams, modalityPinOverride: true }).modality_pin_override).toBe(true);
expect(buildComplexityRouterConfig(baseParams).modality_pin_override).toBe(false);
expect(buildComplexityRouterConfig({ ...baseParams, modalityPinOverride: false }).modality_pin_override).toBe(
false,
);
});
it("writes session_affinity=true so turning the toggle on overrides the backend's off-by-default", () => {
const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: true });
expect(config.session_affinity).toBe(true);

4
ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
View file

@ -111,6 +111,7 @@ export interface BuildComplexityRouterConfigParams {
classificationMode: ClassificationMode | undefined;
sessionAffinity: boolean;
modalityRouting?: boolean;
modalityPinOverride?: boolean;
deploymentAffinity: boolean;
customTechnicalKeywords: string[];
keywordTierRules: KeywordTierRule[];
@ -169,6 +170,7 @@ export interface ComplexityRouterConfigPayload {
session_affinity: boolean;
deployment_affinity: boolean;
modality_routing: boolean;
modality_pin_override: boolean;
custom_technical_keywords?: string[];
keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[];
semantic_keyword_matching?: boolean;
@ -409,6 +411,7 @@ export const buildComplexityRouterConfig = ({
classificationMode,
sessionAffinity,
modalityRouting,
modalityPinOverride,
deploymentAffinity,
customTechnicalKeywords,
keywordTierRules,
@ -472,6 +475,7 @@ export const buildComplexityRouterConfig = ({
session_affinity: sessionAffinity,
deployment_affinity: deploymentAffinity,
modality_routing: modalityRouting ?? false,
modality_pin_override: modalityPinOverride ?? false,
...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }),
...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }),
escalation_keywords: cleanedEscalationKeywords,

26
ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
View file

@ -257,6 +257,30 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => {
});
});
describe("buildUpdatedComplexityRouterConfig modality pin override", () => {
it("writes modality_pin_override explicitly both ways", () => {
expect(
buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, modality_pin_override: true }).modality_pin_override,
).toBe(true);
expect(
buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, modality_pin_override: false }).modality_pin_override,
).toBe(false);
});
it("re-asserts the backend's off-by-default when the form value is absent, rather than dropping the key", () => {
const result = buildUpdatedComplexityRouterConfig({ ...STORED, modality_pin_override: true }, FORM_VALUE);
expect(result.modality_pin_override).toBe(false);
});
it("round-trips a stored modality_pin_override=true through hydrate then save", () => {
const stored = { ...STORED, modality_routing: true, modality_pin_override: true };
const hydrated = hydrateComplexityRouterConfig(stored, undefined);
expect(hydrated.modality_pin_override).toBe(true);
expect(buildUpdatedComplexityRouterConfig(stored, hydrated).modality_pin_override).toBe(true);
});
});
describe("buildUpdatedComplexityRouterConfig classification mode", () => {
it("round-trips a stored user_turn through hydrate then save", () => {
const stored = { ...STORED, classification_mode: "user_turn" };
@ -505,6 +529,8 @@ describe("managed keys survive an untouched open-and-save", () => {
classifier_fallback: "default_model",
classification_mode: "user_turn",
session_affinity: true,
modality_routing: true,
modality_pin_override: true,
deployment_affinity: false,
adaptive: true,
adaptive_weights: { quality: 0.4, cost: 0.6 },

22
ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts
View file

@ -51,6 +51,7 @@ const expectedClassifiedTierConfig = {
session_affinity: false,
deployment_affinity: true,
modality_routing: false,
modality_pin_override: false,
adaptive: true,
adaptive_weights: { quality: 0.4, cost: 0.6 },
adaptive_eligible: "classified_tier",
@ -74,6 +75,7 @@ const expectedAdaptiveDisabledConfig = {
session_affinity: false,
deployment_affinity: true,
modality_routing: false,
modality_pin_override: false,
};
describe("buildUpdatedComplexityRouterConfig", () => {
@ -109,6 +111,26 @@ describe("buildUpdatedComplexityRouterConfig", () => {
expect(disabled.modality_routing).toBe(false);
});
it("hydrates a stored modality_pin_override into form state and defaults absent to off", () => {
expect(
hydrateComplexityRouterConfig({ ...storedConfig, modality_pin_override: true }, null).modality_pin_override,
).toBe(true);
expect(hydrateComplexityRouterConfig(storedConfig, null).modality_pin_override).toBe(false);
});
it("round-trips modality_pin_override explicitly in both directions", () => {
const enabled = buildUpdatedComplexityRouterConfig(storedConfig, {
...classifiedTierValue,
modality_pin_override: true,
});
expect(enabled.modality_pin_override).toBe(true);
const disabled = buildUpdatedComplexityRouterConfig(
{ ...storedConfig, modality_pin_override: true },
{ ...classifiedTierValue, modality_pin_override: false },
);
expect(disabled.modality_pin_override).toBe(false);
});
it("includes return_raw_model_name only when enabled", () => {
const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, {
...classifiedTierValue,

44
ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
View file

@ -549,6 +549,50 @@ describe("EditAutoRouterModal deployment affinity", () => {
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().deployment_affinity).toBe(false);
});
// modality_pin_override is a managed key, so the modal rewrites it from form state on save. A
// hydration gap would silently turn a stored override off on the next untouched save.
it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => {
const user = userEvent.setup();
renderWithStoredConfig({ ...STORED_CONFIG, modality_routing: true, modality_pin_override: true });
await user.click(await screen.findByText("Advanced: Modality Routing"));
expect(await screen.findByRole("switch", { name: "Override session pin for image requests" })).toBeChecked();
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().modality_pin_override).toBe(true);
});
it("persists turning the modality pin override on", async () => {
const user = userEvent.setup();
renderWithStoredConfig({ ...STORED_CONFIG, modality_routing: true });
await user.click(await screen.findByText("Advanced: Modality Routing"));
await user.click(await screen.findByRole("switch", { name: "Override session pin for image requests" }));
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().modality_pin_override).toBe(true);
});
it("writes modality_pin_override=false for a stored config that never carried the key", async () => {
const user = userEvent.setup();
renderWithStoredConfig(STORED_CONFIG);
await user.click(await screen.findByText("Advanced: Modality Routing"));
expect(await screen.findByRole("switch", { name: "Override session pin for image requests" })).toHaveAttribute(
"aria-disabled",
"true",
);
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().modality_pin_override).toBe(false);
});
});
describe("EditAutoRouterModal custom classifier prompt and fallback", () => {

5
ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
View file

@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig {
reasoning_override_min_score?: unknown;
session_affinity?: unknown;
modality_routing?: unknown;
modality_pin_override?: unknown;
deployment_affinity?: unknown;
adaptive?: boolean;
adaptive_weights?: AdaptiveRouterWeights;
@ -181,6 +182,8 @@ export const hydrateComplexityRouterConfig = (
session_affinity:
typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY,
modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false,
modality_pin_override:
typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false,
deployment_affinity:
typeof parsedConfig.deployment_affinity === "boolean"
? parsedConfig.deployment_affinity
@ -221,6 +224,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
"classification_mode",
"session_affinity",
"modality_routing",
"modality_pin_override",
"deployment_affinity",
"adaptive",
"adaptive_weights",
@ -318,6 +322,7 @@ export const buildUpdatedComplexityRouterConfig = (
classifierFallback: value.classifier_fallback,
sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY,
modalityRouting: value.modality_routing ?? false,
modalityPinOverride: value.modality_pin_override ?? false,
deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
customTechnicalKeywords: customTechnicalKeywords ?? [],
keywordTierRules: keywordMatching?.keywordTierRules ?? [],

6
ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx
View file

@ -186,6 +186,12 @@ describe("RoutingDecisionCard", () => {
expect(screen.queryByText("housekeeping")).not.toBeInTheDocument();
});
it("labels a modality pin override instead of showing the raw cause token", () => {
render(<RoutingDecisionCard decision={{ ...heuristic, cause: "modality_pin_override" }} />);
expect(screen.getByText("Overrode session pin for image input")).toBeInTheDocument();
expect(screen.queryByText("modality_pin_override")).not.toBeInTheDocument();
});
it("labels a modality escalation instead of showing the raw cause token", () => {
render(<RoutingDecisionCard decision={{ ...heuristic, cause: "modality_escalation" }} />);
expect(screen.getByText("Escalated for image input")).toBeInTheDocument();

1
ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx
View file

@ -93,6 +93,7 @@ const CONSTANT_CAUSE_LABELS: Record<string, string> = {
session_affinity_escalation: "Escalated from session pin",
user_turn_continuation: "Continuation turn, classifier skipped",
modality_escalation: "Escalated for image input",
modality_pin_override: "Overrode session pin for image input",
quality_tier: "Quality tier mapping",
bandit: "Adaptive bandit",
default_fallback: "Default model, no route matched",

19
ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
View file

@ -147,6 +147,25 @@ describe("autorouter_presets", () => {
expect(withoutFlag.complexityRouterConfig.modality_routing).toBe(false);
});
it("carries a preset's modality_pin_override into the prefilled form state", () => {
const preset = getPresetByKey("anthropic_family")!;
const withFlag = { ...preset.complexity_router_config, modality_routing: true, modality_pin_override: true };
const prefill = buildPresetPrefill(withFlag, groupsOnly(getRequiredModelsInPreset(preset)));
expect(prefill.complexityRouterConfig.modality_pin_override).toBe(true);
const withoutFlag = buildPresetPrefill(
preset.complexity_router_config,
groupsOnly(getRequiredModelsInPreset(preset)),
);
expect(withoutFlag.complexityRouterConfig.modality_pin_override).toBe(false);
});
it("ships every bundled preset with both modality flags written out, since the payload type requires them", () => {
for (const preset of getAllPresets()) {
expect(preset.complexity_router_config.modality_routing, preset.key).toBe(false);
expect(preset.complexity_router_config.modality_pin_override, preset.key).toBe(false);
}
});
it("prefills the anthropic preset's effort through to tier_model_params", () => {
const preset = getPresetByKey("anthropic_family")!;
const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset)));

1
ui/litellm-dashboard/src/lib/autorouter_presets.ts
View file

@ -286,6 +286,7 @@ export const buildPresetPrefill = (
session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
modality_routing: config.modality_routing ?? false,
modality_pin_override: config.modality_pin_override ?? false,
adaptive: config.adaptive,
adaptive_weights: config.adaptive_weights,
tier_distance_penalty: config.tier_distance_penalty,

15
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -25772,6 +25772,11 @@ export interface components {
* @description Number of trusted reverse proxies/load balancers in front of the gateway that append to X-Forwarded-For. When set (and mcp_trusted_proxy_ranges validates the direct peer), the client IP for MCP access control is read this many entries from the right of the chain instead of the spoofable leftmost value, defeating append-style X-Forwarded-For forgery.
*/
mcp_xff_num_trusted_hops?: number | null;
/**
* Missing Session Id
* @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.
*/
missing_session_id?: ("generate" | "reject") | null;
/**
* Model List Healthy Only
* @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called.
@ -34735,9 +34740,15 @@ export interface components {
* @default 0.5
*/
match_threshold: number;
/**
* Modality Pin Override
* @description Let modality_routing replace a kept session-affinity pin on the turns that carry an image. Without this, a session pinned to a text-only model fails every image turn with a provider 400, since the pin is exempt from the modality gate. When enabled, such a turn routes to a capable model for that request only and the stored pin is left untouched, so the next text turn replays the session's own model; the override is reported as cause modality_pin_override and is never itself pinned. Inert unless modality_routing is also enabled.
* @default false
*/
modality_pin_override: boolean;
/**
* Modality Routing
* @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, a routed model explicitly declared supports_vision false (deployment model_info or the model cost map; unmapped names stay routable) is replaced by the nearest HIGHER tier holding a capable model, then default_model, else a clear 400. A kept session-affinity pin still wins even when an image arrives.
* @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, a routed model explicitly declared supports_vision false (deployment model_info or the model cost map; unmapped names stay routable) is replaced by the nearest HIGHER tier holding a capable model, then default_model, else a clear 400. A kept session-affinity pin still wins even when an image arrives, unless modality_pin_override is also enabled.
* @default false
*/
modality_routing: boolean;
@ -35940,7 +35951,7 @@ export interface components {
* Cause
* @enum {string}
*/
cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit";
cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "modality_pin_override" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit";
/** Classifier Cost */
classifier_cost?: number;
/** Classifier Model */