diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 3f96531cf6f..2967fc2a505 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -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 diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 9648321d7ff..0125886aac1 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -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(¶m)); + } + } + + #[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(¶m)); + } + } + + #[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] diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index cb7175691df..16ef6920114 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -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): """ diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 5e116b7301a..ec86c0ae1d9 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -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") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1f552ff3e13..1434fa5bfea 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index af3ff6714af..bb075220530 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index 50c05daee11..28220f09f00 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index d187b5b12e9..7d88a037f4f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -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: diff --git a/litellm/proxy/public_endpoints/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json index 1e0d40567d7..c2b13b81542 100644 --- a/litellm/proxy/public_endpoints/autorouter_presets.json +++ b/litellm/proxy/public_endpoints/autorouter_presets.json @@ -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 } } diff --git a/litellm/router.py b/litellm/router.py index 0af514fe8a2..303b22c9484 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 = [] diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index ad8b67d5e8f..ee51add1ca1 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -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:` 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 diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 430efe339a2..a00ae6bee80 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -751,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") @@ -760,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") ) @@ -2393,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 @@ -2407,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) ): @@ -2449,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", @@ -2458,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), diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 0ae0db63fad..9f2054dda01 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -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." ), ) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index f7855cb38ff..601b32c4386 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -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, diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index bae876dfdd9..17cf5831c96 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -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): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ee6f09e05dc..569fce4f7b8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 98045725177..6b162ea90f1 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -370,6 +370,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( custom_llm_provider, litellm_params, logging_obj, + embedding_executor=None, extra_headers=None, extra_body=None, timeout=None, diff --git a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py index 0c4d1dfc21e..4812d199c06 100644 --- a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py +++ b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py @@ -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"]}]) diff --git a/tests/rust-python-harness/README.md b/tests/rust-python-harness/README.md index e94ac87c3b3..a34e1a1ab73 100644 --- a/tests/rust-python-harness/README.md +++ b/tests/rust-python-harness/README.md @@ -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: diff --git a/tests/rust-python-harness/cli.py b/tests/rust-python-harness/cli.py index 41c46f9613a..c996b68846c 100644 --- a/tests/rust-python-harness/cli.py +++ b/tests/rust-python-harness/cli.py @@ -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) diff --git a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json index abeea01d9b5..d838486772d 100644 --- a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json +++ b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json @@ -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."} } } diff --git a/tests/rust-python-harness/existing_e2e_test_sdk/README.md b/tests/rust-python-harness/existing_e2e_test_sdk/README.md new file mode 100644 index 00000000000..fb84f170703 --- /dev/null +++ b/tests/rust-python-harness/existing_e2e_test_sdk/README.md @@ -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. diff --git a/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json b/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json new file mode 100644 index 00000000000..eefceea1a75 --- /dev/null +++ b/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json @@ -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."} + } +} diff --git a/tests/rust-python-harness/models.py b/tests/rust-python-harness/models.py index 21097e0f7d0..a02684dc63d 100644 --- a/tests/rust-python-harness/models.py +++ b/tests/rust-python-harness/models.py @@ -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) diff --git a/tests/rust-python-harness/runner.py b/tests/rust-python-harness/runner.py index 82393ef234e..286f2c4116d 100644 --- a/tests/rust-python-harness/runner.py +++ b/tests/rust-python-harness/runner.py @@ -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}[" diff --git a/tests/rust-python-harness/shared/__init__.py b/tests/rust-python-harness/shared/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/shared/parity/__init__.py b/tests/rust-python-harness/shared/parity/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/shared/parity/ledger.py b/tests/rust-python-harness/shared/parity/ledger.py new file mode 100644 index 00000000000..40dfed583ae --- /dev/null +++ b/tests/rust-python-harness/shared/parity/ledger.py @@ -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, + ) diff --git a/tests/rust-python-harness/strategies/__init__.py b/tests/rust-python-harness/strategies/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/strategies/unit_tests/__init__.py b/tests/rust-python-harness/strategies/unit_tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json new file mode 100644 index 00000000000..e617ddf8f94 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json @@ -0,0 +1,1084 @@ +{ + "sdk_function": "ocr", + "python_scope": [ + "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "tests/test_litellm/ocr/test_rust_bridge.py", + "tests/test_litellm/ocr/test_ocr_file_input.py", + "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "tests/test_litellm/ocr/test_ocr_native_format.py", + "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py" + ], + "rust_scope": [ + "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs" + ], + "entries": [ + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_should_encode_azure_document_intelligence_model_id", + "status": "unmapped", + "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", + "status": "unmapped", + "reason": "model-id dot-segment validation has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_preserves_azure_native_fields", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_response_normalizes_pages", + "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_response_normalizes_pages", + "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", + "status": "unmapped", + "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_non_succeeded_status_raises", + "status": "unmapped", + "reason": "no Rust test asserts on a non-succeeded Azure DI status" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_supported_ocr_params_includes_features", + "status": "unmapped", + "reason": "supported-params list content has no Rust equivalent for Azure" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", + "status": "unmapped", + "reason": "native req_format raw-operation passthrough is not tested in Rust" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", + "status": "unmapped", + "reason": "native req_format raw-operation passthrough is not tested in Rust" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", + "status": "unmapped", + "reason": "req_format gating of raw-operation output has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_passes_through_req_format", + "status": "unmapped", + "reason": "req_format passthrough in map_ocr_params has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", + "status": "unmapped", + "reason": "req_format validation error path has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_omits_req_format_query_param", + "status": "unmapped", + "reason": "no Rust test asserts req_format is excluded from the built URL" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_features", + "status": "unmapped", + "reason": "features-string normalization in map_ocr_params has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_empty_features_list_omitted", + "status": "unmapped", + "reason": "empty-features omission has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_invalid_features_raises", + "status": "unmapped", + "reason": "features validation error path has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_appends_features_query", + "status": "unmapped", + "reason": "features query-param construction has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_combines_pages_and_features", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_url_normalizes_zero_based_pages", + "justification": "both assert 0-based, duplicate page indices are deduped, sorted, and rewritten 1-based into the request URL" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_validate_environment_uses_subscription_key", + "status": "unmapped", + "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_validate_environment_falls_back_to_entra_token", + "status": "unmapped", + "reason": "Entra bearer-token fallback logic has no Rust test" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", + "status": "unmapped", + "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", + "status": "unmapped", + "reason": "api_base precedence resolution is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", + "status": "unmapped", + "reason": "api_base precedence resolution is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_use_litellm_rust_toggles_flag", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_env_var_enables_rust_ocr", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side env-var flag gating" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_returns_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", + "status": "unmapped", + "reason": "bridge-plumbing: native-extension import/loader fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_loader_caches_absent_extension", + "status": "unmapped", + "reason": "bridge-plumbing: loader caching behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_available_reflects_loader", + "status": "unmapped", + "reason": "bridge-plumbing: loader availability check" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_aocr_returns_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: dependency-injection test hook" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: injected-impl state retention regression" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_explicit_ocr_none_clears_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: injected-impl clearing behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_none_when_extension_absent", + "status": "unmapped", + "reason": "bridge-plumbing: degrade path when the native extension is missing" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_uses_compiled_extension", + "status": "unmapped", + "reason": "bridge-plumbing: native module resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side timeout normalization helper" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: async wrapper argument forwarding" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: request preparation and response wrapping in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", + "status": "unmapped", + "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", + "status": "unmapped", + "reason": "secret-manager: key precedence resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", + "status": "unmapped", + "reason": "secret-manager: provider-specific env var name resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", + "status": "unmapped", + "reason": "secret-manager: vertex routing metadata merge happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", + "status": "unmapped", + "reason": "secret-manager: vertex project/location resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", + "status": "unmapped", + "reason": "secret-manager: azure_ai api_base resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", + "status": "unmapped", + "reason": "secret-manager: doc-intelligence endpoint resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_runs_pre_call_logging", + "status": "unmapped", + "reason": "bridge-plumbing: Python logging-object pre_call invocation" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_routes_to_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: provider-prefix stripping before routing" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", + "status": "unmapped", + "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_exception_type_uses_resolved_provider_context", + "status": "unmapped", + "reason": "bridge-plumbing: Python exception-type mapping on bridge failure" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_aocr_routes_to_async_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: async routing to a fake bridge" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_aocr_exception_type_uses_resolved_provider_context", + "status": "unmapped", + "reason": "bridge-plumbing: async exception-type mapping on bridge failure" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_does_not_route_to_rust_when_disabled", + "status": "unmapped", + "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", + "status": "unmapped", + "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_forwards_timeout_to_rust", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_passes_default_request_timeout_to_rust", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", + "status": "unmapped", + "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_png_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_be_case_insensitive", + "status": "unmapped", + "reason": "file-normalization: MIME detection case handling" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", + "status": "unmapped", + "reason": "file-normalization: MIME detection fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", + "status": "unmapped", + "reason": "file-normalization: arbitrary-file-read guard on bare str paths" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", + "status": "unmapped", + "reason": "file-normalization: raw-bytes-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override on raw bytes" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override on raw bytes" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", + "status": "unmapped", + "reason": "file-normalization: file-like-object conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", + "status": "unmapped", + "reason": "file-normalization: file-like-object name-based MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", + "status": "unmapped", + "reason": "file-normalization: missing-field validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: missing-file validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", + "status": "unmapped", + "reason": "file-normalization: empty-file validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", + "status": "unmapped", + "reason": "file-normalization: unsupported input type validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME-type injection validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override precedence" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", + "status": "unmapped", + "reason": "file-normalization: filename-based MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", + "status": "unmapped", + "reason": "file-normalization: filename-based MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", + "status": "unmapped", + "reason": "file-normalization: default MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", + "status": "unmapped", + "reason": "file-normalization: binary round-trip through base64" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", + "status": "unmapped", + "reason": "file-normalization: content-type parameter stripping" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", + "status": "unmapped", + "reason": "file-normalization: content-type parameter stripping" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body parsing" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body parsing error path" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", + "status": "unmapped", + "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "extract_header_is_a_supported_ocr_param", + "justification": "both assert extract_header appears in the Mistral supported OCR params list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "extract_footer_is_a_supported_ocr_param", + "justification": "both assert extract_footer appears in the Mistral supported OCR params list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "existing_ocr_params_remain_supported", + "justification": "both assert the previously supported params are still present in the supported list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_header_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_header", + "justification": "both assert extract_header alone survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_footer_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_footer", + "justification": "both assert extract_footer alone survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_header_and_footer", + "justification": "both assert header and footer passed together are both forwarded with their given values" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_drops_unknown_params", + "justification": "both assert an unrecognized param key is dropped while a known one is kept" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "new_ocr_params_are_supported", + "justification": "both assert each OCR4 param is in the supported list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_new_ocr_params", + "justification": "both assert each OCR4 param/value pair survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_includes_each_optional_param", + "justification": "both assert each optional param value lands in the built request body alongside model/document with no files" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_includes_multiple_new_params", + "justification": "both assert multiple OCR4 params passed together all land in the same request body" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_preserves_blocks_and_confidence_scores", + "justification": "both assert blocks and confidence_scores survive the OCR response transform on the returned page" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_preserves_ocr4_page_fields", + "justification": "both assert tables, hyperlinks, header and footer survive the OCR response transform on the returned page" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_model_info_ocr4_price", + "status": "unmapped", + "reason": "cost-calc: pricing/model-info lookup is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr4_cost_scales_with_pages", + "status": "unmapped", + "reason": "cost-calc: per-page pricing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_pricing_entry", + "status": "unmapped", + "reason": "cost-calc: cost-map JSON entry validation is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_model_info_price", + "status": "unmapped", + "reason": "cost-calc: pricing/model-info lookup is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_cost_scales_with_pages", + "status": "unmapped", + "reason": "cost-calc: per-page pricing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", + "status": "unmapped", + "reason": "cost-calc: mixed-rate billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_annotation_only_response", + "status": "unmapped", + "reason": "cost-calc: annotation-only billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", + "status": "unmapped", + "reason": "cost-calc: fallback billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", + "status": "unmapped", + "reason": "cost-calc: fallback billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_rust_ocr_serves_default_format", + "status": "unmapped", + "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_rust_ocr_skipped_for_native_format", + "status": "unmapped", + "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", + "status": "unmapped", + "reason": "provider-support validation for req_format happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", + "status": "unmapped", + "reason": "req_format validation error path is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", + "status": "unmapped", + "reason": "guardrail-translation handler discovery is a Python proxy-layer concern" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", + "status": "unmapped", + "reason": "guardrail-translation handler discovery is a Python proxy-layer concern" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_document_url", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_image_url", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_no_document", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_invalid_document", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_input_blocking_guardrail", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_single_page", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_multiple_pages", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_empty_pages", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_output_blocking_guardrail", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", + "status": "unmapped", + "reason": "PII redaction in the translation handler has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_read_req_format_from_header", + "status": "unmapped", + "reason": "proxy-layer header parsing has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_prefer_body_req_format_over_header", + "status": "unmapped", + "reason": "proxy-layer body-vs-header precedence has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_omit_req_format_when_header_absent", + "status": "unmapped", + "reason": "proxy-layer parsing has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_reject_unknown_req_format", + "status": "unmapped", + "reason": "proxy-layer validation has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_return_native_payload_with_litellm_response_headers", + "status": "unmapped", + "reason": "proxy-layer response construction has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_return_normalized_response_when_no_native_payload", + "status": "unmapped", + "reason": "proxy-layer response construction has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_explicit_false_overrides_process_enable", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the explicit-per-request rust:False override wins over the Python process flag, no Rust-owned behavior runs" + } + ], + "rust_only_tests": [ + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", + "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "no_callback_fast_path_dispatches_nothing", + "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", + "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "blocks_private_and_metadata_ips", + "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "convert_document_url_rejects_loopback_fetch", + "reason": "URL-fetch SSRF protection is Rust-gateway-only" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "convert_document_url_leaves_data_uri_untouched", + "reason": "URL-fetch SSRF protection is Rust-gateway-only" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_passes_short_strings_through", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_caps_long_payloads", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_does_not_split_multibyte_chars", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_dispatch_supports_migrated_providers", + "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "string_headers_accepts_string_values", + "reason": "Rust-gateway header-coercion helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "auth_header_detection_is_case_insensitive", + "reason": "Rust-gateway header-detection helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", + "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", + "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", + "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", + "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", + "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "string_headers_rejects_non_string_values", + "reason": "Rust-gateway header-coercion error path has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "azure_ai_reuses_mistral_body_transform", + "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", + "reason": "no Python test asserts on the base64Source request body shape" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_mistral_url_uses_project_location_and_model", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_mistral_reuses_mistral_body_transform", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_deepseek_response_wraps_markdown_content", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_rejects_non_object_document", + "reason": "non-object document rejection has no dedicated Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_normalizes_mistral_json", + "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "complete_url_defaults_and_dedupes_v1", + "reason": "URL-building/defaulting for Mistral has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "resolve_api_key_prefers_param_then_env", + "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "resolve_api_key_errors_when_absent", + "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", + "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer" + } + ] +} diff --git a/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py new file mode 100644 index 00000000000..b8355e04266 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py @@ -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) + ) diff --git a/tests/rust-python-harness/strategies/unit_tests/python_runner.py b/tests/rust-python-harness/strategies/unit_tests/python_runner.py new file mode 100644 index 00000000000..a7528d27756 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/python_runner.py @@ -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) diff --git a/tests/rust-python-harness/strategies/unit_tests/rust_runner.py b/tests/rust-python-harness/strategies/unit_tests/rust_runner.py new file mode 100644 index 00000000000..6b855adbc4b --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/rust_runner.py @@ -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)) diff --git a/tests/rust-python-harness/ui.py b/tests/rust-python-harness/ui.py index 57fedd17fa6..bf143a4e100 100644 --- a/tests/rust-python-harness/ui.py +++ b/tests/rust-python-harness/ui.py @@ -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 = [] diff --git a/tests/rust-python-harness/unit_tests_rust/strategy.json b/tests/rust-python-harness/unit_tests_rust/strategy.json index 89e897c872d..bfb2bc0dad0 100644 --- a/tests/rust-python-harness/unit_tests_rust/strategy.json +++ b/tests/rust-python-harness/unit_tests_rust/strategy.json @@ -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": []} } } diff --git a/tests/rust-python-harness/validate_sub_methods/strategy.json b/tests/rust-python-harness/validate_sub_methods/strategy.json index 6e6381678e0..a26bc2d70ed 100644 --- a/tests/rust-python-harness/validate_sub_methods/strategy.json +++ b/tests/rust-python-harness/validate_sub_methods/strategy.json @@ -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": []} } } diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py index 16c518ff412..a96eae0f9c3 100644 --- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py +++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py @@ -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() diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py index 2d0605e3b7f..34c62864c4e 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -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( diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py index 7e1139ca79a..222ead92b54 100644 --- a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py +++ b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py @@ -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( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 3f6d8f8837c..ff0007f47aa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5508259273d..5cde2e83f62 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index 1b50ea53db2..a07157396df 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index da66c36328e..47089b7b1b1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -15,6 +15,7 @@ import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( @@ -3778,3 +3779,1153 @@ async def test_moderation_hook_skips_chat_traffic_when_configured_for_during_mcp assert result == data mock_post.assert_not_called() + + +_ANTHROPIC_SSE_CHUNKS = ( + b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message",' + b'"role":"assistant","model":"claude","content":[],"usage":{"input_tokens":5,"output_tokens":0}}}\n\n', + b'event: content_block_start\ndata: {"type":"content_block_start","index":0,' + b'"content_block":{"type":"text","text":""}}\n\n', + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"my card is 4111-1111-1111-1111"}}\n\n', + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', + b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"output_tokens":9}}\n\n', + b'event: message_stop\ndata: {"type":"message_stop"}\n\n', +) + +_MODEL_ARMOR_CLEAN = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND"}} + +_MODEL_ARMOR_BLOCK = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": [ + {"infoType": "CREDIT_CARD_NUMBER", "likelihood": "VERY_LIKELY"} + ], + } + } + } + }, + } +} + +# The root-level sanitizedText fallback in _get_sanitized_content, i.e. a rewrite that trips no +# named filter +_MODEL_ARMOR_SANITIZED = { + "sanitizedText": "my card is [REDACTED]", + "sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND"}, +} + +# The shape a real de-identify template returns: the SDP filter both matches and hands back the +# rewritten text, so whether it blocks or masks is decided by allow_sanitization alone +_MODEL_ARMOR_DEIDENTIFIED = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": { + "matchState": "MATCH_FOUND", + "data": {"text": "my card is [REDACTED]"}, + } + } + } + }, + } +} + + +def _chat_completion_chunks(): + """The chat-completions surface: typed ModelResponseStream chunks.""" + return ( + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content="my card is 4111-1111-1111-1111"), + ) + ] + ), + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content=""), + finish_reason="stop", + ) + ] + ), + ) + + +def _surface_guardrail(**kwargs): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + **kwargs, + ) + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + return guardrail + + +def _armor_post_mock(payload): + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value=payload) + return AsyncMock(return_value=mock_response) + + +async def _anthropic_sse_stream(): + for chunk in _ANTHROPIC_SSE_CHUNKS: + yield chunk + + +def _responses_api_events(): + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + completed = ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + return ( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="my card is 4111-1111-1111-1111", + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=completed, + ), + ) + + +async def _drain_surface_hook(guardrail, chunks, request_data=None): + async def _stream(): + for chunk in chunks: + yield chunk + + return [ + item + async for item in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=_stream(), + request_data=request_data + if request_data is not None + else { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "show me a card"}], + "metadata": {"guardrails": ["model-armor-test"]}, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_streaming_hook_scans_raw_anthropic_sse_instead_of_crashing(): + """A /v1/messages stream arrives as raw SSE bytes and must be assembled, then scanned. + + Regression for the 500 `Error building chunks for logging/streaming usage calculation`: + stream_chunk_builder calls .get() on each chunk, which raises on bytes. + """ + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, _ANTHROPIC_SSE_CHUNKS) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "my card is 4111-1111-1111-1111" in scanned + assert tuple(delivered) == _ANTHROPIC_SSE_CHUNKS + + +@pytest.mark.asyncio +async def test_streaming_hook_scans_responses_api_events_instead_of_crashing(): + """A /v1/responses stream arrives as typed Responses events, which stream_chunk_builder + cannot subscript. The final response.completed event carries the text to scan.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + events = _responses_api_events() + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, events) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert scanned == "my card is 4111-1111-1111-1111" + assert tuple(delivered) == events + + +@pytest.mark.asyncio +async def test_streaming_block_emits_anthropic_error_frame(): + """A block on /v1/messages must terminate the stream in Anthropic's error format. + + The OpenAI-shaped `data: {"error": ...}` frame the chat surface uses is rejected by + Anthropic clients. + """ + guardrail = _surface_guardrail() + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_BLOCK) + ): + delivered = await _drain_surface_hook(guardrail, _ANTHROPIC_SSE_CHUNKS) + + body = b"".join(delivered) + assert b"event: error" in body + assert b'"type": "error"' in body + assert b"guardrail_error" in body + assert b"Streaming response blocked by Model Armor" in body + assert b"4111-1111-1111-1111" not in body + + +@pytest.mark.asyncio +async def test_streaming_block_emits_responses_api_error_event(): + """A block on /v1/responses must terminate the stream with a Responses ErrorEvent.""" + from litellm.types.llms.openai import ErrorEvent + + guardrail = _surface_guardrail() + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_BLOCK) + ): + delivered = await _drain_surface_hook(guardrail, _responses_api_events()) + + assert len(delivered) == 1 + error_event = delivered[0] + assert isinstance(error_event, ErrorEvent) + assert error_event.error.type == "guardrail_error" + assert error_event.error.code == "400" + assert error_event.error.message == "Streaming response blocked by Model Armor" + + +@pytest.mark.asyncio +async def test_streaming_masking_re_emits_anthropic_sse_with_sanitized_text(): + """mask_response_content on /v1/messages must ship the sanitized text, not the original.""" + guardrail = _surface_guardrail(mask_response_content=True) + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) + ): + delivered = await _drain_surface_hook(guardrail, _ANTHROPIC_SSE_CHUNKS) + + body = b"".join(delivered) + assert b"[REDACTED]" in body + assert b"4111-1111-1111-1111" not in body + + +@pytest.mark.asyncio +async def test_streaming_masking_blocks_responses_api_stream(): + """A Responses event stream cannot be rebuilt from sanitized text, so releasing it would + ship the content the guardrail just rewrote. It is blocked instead.""" + from litellm.types.llms.openai import ErrorEvent + + guardrail = _surface_guardrail(mask_response_content=True) + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) + ): + delivered = await _drain_surface_hook(guardrail, _responses_api_events()) + + assert len(delivered) == 1 + assert isinstance(delivered[0], ErrorEvent) + assert delivered[0].error.code == "400" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["anthropic_sse", "responses"]) +async def test_streaming_api_failure_frames_error_per_surface(surface): + """A Model Armor outage with fail_on_error must terminate the stream in the endpoint's + own error format rather than leaking an OpenAI SSE frame onto it.""" + from litellm.types.llms.openai import ErrorEvent + + guardrail = _surface_guardrail(fail_on_error=True) + chunks = _ANTHROPIC_SSE_CHUNKS if surface == "anthropic_sse" else _responses_api_events() + + mock_response = AsyncMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=mock_response) + ): + delivered = await _drain_surface_hook(guardrail, chunks) + + assert len(delivered) >= 1 + if surface == "anthropic_sse": + assert b"event: error" in b"".join(delivered) + else: + assert isinstance(delivered[0], ErrorEvent) + assert delivered[0].error.code == "500" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunks", + [ + pytest.param( + (b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"hi"}}\n\n',), + id="anthropic-sse-without-message-start", + ), + pytest.param(None, id="responses-stream-without-completed-event"), + pytest.param("created", id="responses-stream-cut-off-after-response-created"), + ], +) +async def test_streaming_hook_fails_closed_when_a_surface_stream_cannot_be_assembled(chunks): + """Forwarding an unscannable /v1/messages or /v1/responses stream would silently disable the + guardrail, so the stream is refused in its own wire format instead of released unscanned.""" + from litellm.types.llms.openai import ( + ErrorEvent, + OutputTextDeltaEvent, + ResponsesAPIStreamEvents, + ) + + if chunks is None or chunks == "created": + delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="my card is 4111-1111-1111-1111", + ) + # response.created carries a ResponsesAPIResponse too, but an empty one: reading the body + # off it would scan "" and release every buffered delta unscanned + chunks = (delta,) if chunks is None else (_responses_created_event(), delta) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) != tuple(chunks) + if isinstance(chunks[0], bytes): + joined = b"".join(item.encode() if isinstance(item, str) else item for item in delivered).decode() + assert "event: error" in joined + assert "could not be assembled for scanning" in joined + return + assert len(delivered) == 1 + assert isinstance(delivered[0], ErrorEvent) + assert "could not be assembled for scanning" in delivered[0].error.message + + +@pytest.mark.asyncio +async def test_streaming_hook_forwards_a_preceding_guardrails_error_item(): + """A guardrail earlier in the post_call chain replaces the stream with its own terminal + error item. That item is not a chat delta, and feeding it to stream_chunk_builder is what + surfaced the ticket's 500, so it has to be forwarded untouched instead.""" + from litellm.types.llms.openai import ( + ErrorEvent, + ErrorEventError, + ResponsesAPIStreamEvents, + ) + + chunks = ( + ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=1, + error=ErrorEventError( + type="guardrail_error", + code="400", + message="Streaming response blocked by Model Armor", + param=None, + ), + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) == chunks + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunks", + [ + pytest.param(None, id="anthropic-error-event"), + pytest.param( + ('data: {"error": {"message": "Streaming response blocked by the first guardrail", "code": "400"}}\n\n',), + id="chat-completions-error-payload", + ), + ], +) +async def test_streaming_hook_forwards_a_preceding_guardrails_error_frame(chunks): + """Chained post_call guardrails hand each other their output. An earlier guardrail's error + frame carries no message to assemble, and replacing it would hide the real refusal.""" + if chunks is None: + chunks = anthropic_sse_error_frames("Streaming response blocked by the first guardrail") + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) == chunks + + +@pytest.mark.asyncio +async def test_streaming_responses_error_falls_back_to_sse_when_the_handler_declines(): + """build_stream_error_items may return None, which must not swallow the block into a clean + 200: the refusal falls back to the chat-completions SSE form that still carries the status.""" + from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( + _StreamSurface, + ) + + class _DecliningGuardrail(ModelArmorGuardrail): + @staticmethod + def _build_responses_error_items(exc): + return None + + guardrail = _DecliningGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + exc = HTTPException(status_code=400, detail={"message": "blocked"}) + + items = guardrail._stream_error_items(exc, surface=_StreamSurface.RESPONSES) + + assert len(items) == 1 + assert '"code": "400"' in items[0] + assert "blocked" in items[0] + + +def _responses_created_event(): + from litellm.types.llms.openai import ( + ResponseCreatedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + return ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_refuses_an_opaque_sse_stream_without_anthropic_framing(): + """/v1/messages is not the only endpoint that streams raw SSE: the Google generateContent + route marks its own stream raw too. Its frames carry no Anthropic event types, so refusing + them in Anthropic's format would hand a Google client a body it cannot parse.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = (b'data: {"candidates":[{"content":{"parts":[{"text":"my card is 4111"}]}}]}\n\n',) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) != chunks + body = "".join(item.decode() if isinstance(item, bytes) else item for item in delivered) + assert "could not be assembled for scanning" in body + assert "event: error" not in body + assert '"code": "500"' in body + + +@pytest.mark.asyncio +async def test_streaming_unassemblable_stream_is_forwarded_when_fail_on_error_is_disabled(): + """fail_on_error: false is a deliberate choice to degrade open, and it governs every other + path in this hook. The fail-closed refusal has to honour it too.""" + guardrail = _surface_guardrail(fail_on_error=False) + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = ( + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"hi"}}\n\n', + ) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) == chunks + + +@pytest.mark.asyncio +async def test_streaming_fail_closed_records_the_applied_guardrail(): + """A refusal that no header or log attributes to the guardrail leaves on-call unable to tell + a guardrail block apart from a provider failure.""" + guardrail = _surface_guardrail() + request_data = { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "show me a card"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + chunks = ( + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"hi"}}\n\n', + ) + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_CLEAN)): + await _drain_surface_hook(guardrail, chunks, request_data=request_data) + + assert request_data["metadata"]["applied_guardrails"] == ["model-armor-test"] + + +@pytest.mark.asyncio +async def test_streaming_responses_tool_call_output_is_scanned(): + """An agentic /v1/responses turn can carry its whole payload in tool-call arguments, which + is what the chat surface already folds into the scanned text.""" + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "send_email", + "arguments": '{"body": "my card is 4111-1111-1111-1111"}', + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (completed,)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "4111-1111-1111-1111" in scanned + assert tuple(delivered) == (completed,) + + +@pytest.mark.asyncio +async def test_streaming_hook_refuses_a_content_stream_that_ends_with_an_error_frame(): + """The chain-aware passthrough must stay narrow. A stream carrying real content plus a + trailing error frame is not a bare refusal to forward: the assembler cannot read it, and + releasing it would ship the buffered content unscanned.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = (*_ANTHROPIC_SSE_CHUNKS, *anthropic_sse_error_frames("upstream gave up")) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + body = b"".join(delivered) + assert b"4111-1111-1111-1111" not in body + assert b"could not be assembled for scanning" in body + + +@pytest.mark.parametrize( + "chunks, expected, case", + [ + (anthropic_sse_error_frames("blocked upstream"), True, "anthropic-error-frames-only"), + ((f"data: {json.dumps({'error': {'message': 'blocked'}})}\n\n",), True, "chat-error-payload-only"), + ((), False, "empty-stream"), + ( + (b'event: message_delta\ndata: {"type":"message_delta","error":null}\n\n',), + False, + "content-event-carrying-a-null-error-field", + ), + ( + ( + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content="my card is 4111-1111-1111-1111"), + ) + ] + ), + *anthropic_sse_error_frames("upstream gave up"), + ), + False, + "typed-content-chunks-plus-a-trailing-error-frame", + ), + ], +) +def test_is_sse_error_stream_only_matches_a_stream_that_is_nothing_but_refusals(chunks, expected, case): + """The chain-aware passthrough turns on this predicate, so anything it calls error-only is + forwarded to the client untouched. A stream that still carries content must not qualify: the + frames-only join drops typed chunks, and a content event may carry an empty ``error`` field.""" + from litellm.proxy.guardrails.anthropic_sse import is_sse_error_stream + + assert is_sse_error_stream(chunks) is expected, case + + +@pytest.mark.asyncio +async def test_streaming_hook_does_not_forward_typed_chunks_that_end_with_an_error_frame(): + """A stream mixing buffered content with a trailing refusal is not the bare refusal the chain + passthrough exists for. Forwarding it would release the content no scanner ever saw.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = ( + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content="my card is 4111-1111-1111-1111"), + ) + ] + ), + *anthropic_sse_error_frames("upstream gave up"), + ) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + body = b"".join(item if isinstance(item, bytes) else str(item).encode() for item in delivered) + assert b"4111-1111-1111-1111" not in body + assert b"could not be assembled for scanning" in body + + +def _delivered_bytes(delivered): + return b"".join( + item + if isinstance(item, bytes) + else item.encode() + if isinstance(item, str) + else str(item.model_dump() if hasattr(item, "model_dump") else item).encode() + for item in delivered + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks, case", [(None, "chat_completions"), (_ANTHROPIC_SSE_CHUNKS, "anthropic_sse")]) +async def test_streaming_deidentify_match_masks_when_masking_is_enabled(chunks, case): + """A de-identify template reports MATCH_FOUND for every redaction it makes, so reading that + match as a refusal makes mask_response_content unusable on a stream: the client gets an error + where its non-streaming sibling gets redacted text. The block check has to allow sanitization + exactly as the non-streaming hook does.""" + guardrail = _surface_guardrail(mask_response_content=True) + post = _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook( + guardrail, _chat_completion_chunks() if chunks is None else chunks + ) + + body = _delivered_bytes(delivered) + assert b"[REDACTED]" in body, case + assert b"4111-1111-1111-1111" not in body, case + assert b"blocked by Model Armor" not in body, case + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks, case", [(None, "chat_completions"), (_ANTHROPIC_SSE_CHUNKS, "anthropic_sse")]) +async def test_streaming_deidentify_match_still_blocks_when_masking_is_disabled(chunks, case): + """Without mask_response_content there is nowhere to put the rewritten text, so the same + de-identify match must still end the stream rather than release the original.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook( + guardrail, _chat_completion_chunks() if chunks is None else chunks + ) + + body = _delivered_bytes(delivered) + assert b"Streaming response blocked by Model Armor" in body, case + assert b"4111-1111-1111-1111" not in body, case + + +@pytest.mark.asyncio +async def test_streaming_deidentify_match_logs_masked_run_as_success_not_blocked(): + """The status stamped on request metadata feeds the spend log, so it has to agree with what + the client actually received: a masked stream is a success, not a block.""" + guardrail = _surface_guardrail(mask_response_content=True) + request_data = { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "show me a card"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): + await _drain_surface_hook(guardrail, _chat_completion_chunks(), request_data=request_data) + + assert request_data["metadata"]["_model_armor_status"] == "success" + + +# A de-identify template that matched but handed back no rewrite, e.g. because the transformation +# itself failed. The match still says the buffered original carries what it matched on +_MODEL_ARMOR_DEIDENTIFIED_NO_TEXT = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": {"sdpFilterResult": {"deidentifyResult": {"matchState": "MATCH_FOUND"}}} + }, + } +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks, case", [(None, "chat_completions"), (_ANTHROPIC_SSE_CHUNKS, "anthropic_sse")]) +async def test_streaming_deidentify_match_without_a_rewrite_fails_closed(chunks, case): + """Allowing sanitization past the block check is a promise to apply the redaction. When Model + Armor matches but returns no sanitized text there is nothing to apply, and yielding the + buffered chunks would hand back exactly what it matched on.""" + guardrail = _surface_guardrail(mask_response_content=True) + post = _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED_NO_TEXT) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook( + guardrail, _chat_completion_chunks() if chunks is None else chunks + ) + + body = _delivered_bytes(delivered) + assert b"4111-1111-1111-1111" not in body, case + assert b"Streaming response blocked by Model Armor" in body, case + + +@pytest.mark.asyncio +async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_as_blocked(): + """The Responses surface has no assembled body to rewrite, so a de-identify match ends as a + refusal. The status stamped on metadata feeds the spend log and has to say so rather than + reporting the success the block check alone would have implied.""" + guardrail = _surface_guardrail(mask_response_content=True) + request_data = { + "model": "gpt-4o-mini", + "input": "show me a card", + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): + delivered = await _drain_surface_hook( + guardrail, _responses_api_events(), request_data=request_data + ) + + body = _delivered_bytes(delivered) + assert b"4111-1111-1111-1111" not in body + assert b"Streaming response blocked by Model Armor" in body + assert request_data["metadata"]["_model_armor_status"] == "blocked" + + +def _responses_api_events_truncated(terminal: str): + """A /v1/responses stream whose text went out as deltas and whose terminal event reports no body. + + ``response.failed`` and ``response.incomplete`` are terminal like ``response.completed``, but a + turn that broke mid-generation reports an empty ``output`` while the deltas ahead of it already + spelled the answer out to the client. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + empty_body = ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + terminal_event = ( + ResponseFailedEvent(type=ResponsesAPIStreamEvents.RESPONSE_FAILED, response=empty_body) + if terminal == "failed" + else ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=empty_body) + ) + return ( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="my card is 4111-1111-1111-1111", + ), + terminal_event, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal", ["failed", "incomplete"]) +async def test_streaming_responses_terminal_event_without_a_body_still_scans_the_deltas(terminal): + """A /v1/responses turn that broke mid-generation has still delivered its deltas. + + Reading only the terminal body would find nothing to scan and hand every buffered delta to the + client untouched, so the deltas themselves are what gets scanned. + """ + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, _responses_api_events_truncated(terminal)) + + post.assert_called_once() + assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"] + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_mcp_argument_deltas_are_scanned_when_the_body_is_empty(): + """A turn that only streamed MCP tool arguments still handed the client a payload. + + The delta fallback is read off the event enum rather than listed by hand, so an argument event + that carries no `output_text` cannot fall out of the scan. + """ + from litellm.types.llms.openai import ( + MCPCallArgumentsDeltaEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + chunks = ( + MCPCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, + output_index=0, + item_id="mcp_1", + delta='{"note": "my card is 4111-1111-1111-1111"}', + sequence_number=0, + ), + ResponseIncompleteEvent( + type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_called_once() + assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"] + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "with_output_text_delta", + [True, False], + ids=["summary-and-text-deltas", "summary-delta-only"], +) +async def test_streaming_responses_reasoning_summary_deltas_are_scanned_alongside_the_body(with_output_text_delta): + """A reasoning turn streams its summary in deltas the terminal body never repeats. + + Reading only the body scans the visible answer and hands the client every summary delta + unscanned, so the body and the deltas are scanned together. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ReasoningSummaryTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + answer = "the weather is fine" + summary_delta = ReasoningSummaryTextDeltaEvent( + type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id="rs_1", + output_index=0, + delta="the user said my card is 4111-1111-1111-1111", + ) + text_deltas = ( + ( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=1, + content_index=0, + delta=answer, + ), + ) + if with_output_text_delta + else () + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (summary_delta, *text_deltas, completed)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "4111-1111-1111-1111" in scanned + assert answer in scanned + assert scanned.count(answer) == 1 + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_deltas_of_separate_fields_do_not_form_a_finding_across_their_boundary(): + """Two fields of a turn are separate text, so what runs across their boundary is not model output. + + A reasoning summary ending in half a card number and an answer opening with the other half + each carry nothing to find, and joining them without a break would invent one. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ReasoningSummaryTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + answer = "1111-1111 is not a full card" + summary_delta = ReasoningSummaryTextDeltaEvent( + type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id="rs_1", + output_index=0, + delta="the prefix they gave me is 4111-1111-", + ) + text_delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=1, + content_index=0, + delta=answer, + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (summary_delta, text_delta, completed)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "4111-1111-" in scanned + assert answer in scanned + assert "4111-1111-1111-1111" not in scanned + rendered = "".join(str(item) for item in delivered) + assert "Streaming response blocked by Model Armor" not in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_one_fields_deltas_still_join_into_a_single_finding(): + """A card number split across two deltas of one field is still one card number to scan.""" + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + halves = ("my card is 4111-1111-", "1111-1111") + text_deltas = tuple( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=half, + ) + for half in halves + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed)) + + post.assert_called_once() + assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"] + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_fields_the_body_repeats_are_not_scanned_a_second_time(): + """A turn whose visible fields all reach the terminal body is scanned once, not twice. + + Two output_text fields stream as deltas and come back in the completed body, so scanning the + deltas on top of the body would send Model Armor two copies of everything the client sees. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + paragraphs = ("the first thing to know", "a second and separate point") + text_deltas = tuple( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id=f"msg_{index}", + output_index=index, + content_index=0, + delta=paragraph, + ) + for index, paragraph in enumerate(paragraphs) + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[ + { + "type": "message", + "id": f"msg_{index}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": paragraph, "annotations": []}], + } + for index, paragraph in enumerate(paragraphs) + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert [scanned.count(paragraph) for paragraph in paragraphs] == [1, 1] + rendered = "".join(str(item) for item in delivered) + assert all(paragraph in rendered for paragraph in paragraphs) + + +def test_every_responses_delta_event_is_in_the_scanned_set(): + """Every ``.delta`` the Responses event enum defines is model output on its way to the client.""" + from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( + _RESPONSES_DELTA_EVENT_TYPES, + ) + from litellm.types.llms.openai import ResponsesAPIStreamEvents + + missing = { + event.value + for event in ResponsesAPIStreamEvents + if event.value.endswith(".delta") and event.value not in _RESPONSES_DELTA_EVENT_TYPES + } + assert not missing + assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index d7d02544efb..3ecccb673f3 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10797,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"), ], ) @@ -10847,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 @@ -10859,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 = [] @@ -11025,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"} diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index 90dc38663dd..514446577fd 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -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" + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3c4c7760c6..5c25b8722ef 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -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 diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 349d03356f8..9b647badf69 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -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(); + 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(); + 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( + , + ); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + expect(screen.getByRole("switch", { name: "Override session pin for image requests" })).toBeChecked(); + }); }); describe("ComplexityRouterConfig affinity panel", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 73ab4f0abef..29dd3d9ea83 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx index dd697b35239..8d6c4ed5d5e 100644 --- a/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx @@ -7,20 +7,36 @@ import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; export const ModalityRoutingControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; -}> = ({ value, onChange }) => ( - <> -
- onChange({ ...value, modality_routing: modalityRouting })} - aria-label="Route image requests to vision-capable models" - /> - Route image requests to vision-capable models -
- - 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. - - -); +}> = ({ value, onChange }) => { + const modalityRouting = value.modality_routing ?? false; + return ( + <> +
+ onChange({ ...value, modality_routing: nextModalityRouting })} + aria-label="Route image requests to vision-capable models" + /> + Route image requests to vision-capable models +
+ + 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. + +
+ onChange({ ...value, modality_pin_override: modalityPinOverride })} + disabled={!modalityRouting} + aria-label="Override session pin for image requests" + /> + Override session pin for image requests +
+ + 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. + + + ); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 7769042a0c7..27a169524b9 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -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(); + + 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(); + + 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", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index ab5ea4cbfc9..fc3fdaf39cd 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -358,6 +358,7 @@ const AddAutoRouterTab: React.FC = ({ 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, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index c6370541a45..3ef55b9295d 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -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); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 94c5badf6d1..91e4fdb23b5 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index fae6744d3ac..01198bcc118 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -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 }, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index 481ba2b6b00..651d5880db4 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 96e8549eac8..6067e72e547 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -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", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 7aa5f4beb1f..442e1ec6039 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -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 ?? [], diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx index bf6a20a4af8..fd1777f802c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx @@ -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(); + 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(); expect(screen.getByText("Escalated for image input")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index ffef01010b7..cf2c71e64c6 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -93,6 +93,7 @@ const CONSTANT_CAUSE_LABELS: Record = { 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", diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index e3132067f4c..44ac6973a1c 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -148,6 +148,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))); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 30f8ed99e74..ff482bd23b8 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -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, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 491c4fb6a44..b0f8e618645 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34740,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; @@ -35945,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 */