diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index c112bf2bb22..d02f5878396 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -128,6 +128,9 @@ jobs: - name: check_fastuuid_usage run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py + - name: check_py310_typing_imports + run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py + - name: check_e2e_no_raw_requests run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -145,3 +148,33 @@ jobs: - name: documentation_test_api_docs run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py + + python-310-import-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.10" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: uv sync --frozen --extra proxy --python 3.10 + + - run: uv run --no-sync python --version + + - name: Import litellm + run: uv run --no-sync python -c "import litellm" + + - name: Check litellm CLI + run: uv run --no-sync litellm --version 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/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index f441ce71ab9..1fddc527ec8 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -11,17 +11,35 @@ import sys sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +import functools import tempfile -from typing import Optional +from contextvars import ContextVar +from typing import TYPE_CHECKING, ClassVar, Literal, Optional from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import walk_user_text +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj GUARDRAIL_NAME = "hide_secrets" +GUARDRAIL_PROVIDER = "hide-secrets" + +# Per-invocation tally of redacted secrets by detect-secrets plugin type; None +# means the guardrail did not run, so _process_response records nothing. +_masked_entity_count: ContextVar[Optional[dict]] = ContextVar( + "hide_secrets_masked_entity_count", default=None +) + _custom_plugins_path = "file://" + os.path.join( os.path.dirname(os.path.abspath(__file__)), "secrets_plugins" ) @@ -422,6 +440,10 @@ _default_detect_secrets_config = { class _ENTERPRISE_SecretDetection(CustomGuardrail): + # Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail + # path skips should_run_check and never sees data["prompt"]). + use_native_lifecycle_hooks: ClassVar[bool] = True + def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs): self.user_defined_detect_secrets_config = detect_secrets_config super().__init__(**kwargs) @@ -455,6 +477,26 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): return detected_secrets + def redact_text(self, text: str, source: str = "message") -> str: + """Replace every detected secret in ``text`` with ``[REDACTED]`` and + tally the detected types into the per-invocation masked-entity count.""" + detected_secrets = self.scan_message_for_secrets(text) + if not detected_secrets: + return text + counts = _masked_entity_count.get() + if counts is not None: + for secret in detected_secrets: + counts[secret["type"]] = counts.get(secret["type"], 0) + 1 + secret_types = [secret["type"] for secret in detected_secrets] + verbose_proxy_logger.warning( + f"Detected and redacted secrets in {source}: {secret_types}" + ) + return functools.reduce( + lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"), + detected_secrets, + text, + ) + async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool: if user_api_key_dict.permissions is not None: if GUARDRAIL_NAME in user_api_key_dict.permissions: @@ -463,7 +505,45 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): return True + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Unified-interface entrypoint, used by /guardrails/apply_guardrail + (the UI test playground). Proxied traffic keeps using + ``async_pre_call_hook``, see ``use_native_lifecycle_hooks``.""" + texts = inputs.get("texts") + if not texts or not any(texts): + return inputs + _masked_entity_count.set({}) + return {**inputs, "texts": [self.redact_text(text) for text in texts]} + + def _redact_prompt(self, data: dict) -> int: + """Redact ``data["prompt"]`` (the text-completion shape, which + ``walk_user_text`` does not cover) and return how many non-empty + strings were inspected.""" + prompt = data.get("prompt") + if isinstance(prompt, str): + if not prompt: + return 0 + data["prompt"] = self.redact_text(prompt, source="prompt") + return 1 + if isinstance(prompt, list): + data["prompt"] = [ # mutable-ok: data["prompt"] is a list on the wire + self.redact_text(item, source="prompt") + if isinstance(item, str) and item + else item + for item in prompt + ] + return sum(1 for item in prompt if isinstance(item, str) and item) + return 0 + #### CALL HOOKS - proxy only #### + @log_guardrail_information async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -471,53 +551,84 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): data: dict, call_type: str, # "completion", "embeddings", "image_generation", "moderation" ): + _masked_entity_count.set(None) if await self.should_run_check(user_api_key_dict) is False: return + _masked_entity_count.set({}) + # Covers multimodal list content + Responses-API input. - def _redact_message_text(text: str) -> str: - detected_secrets = self.scan_message_for_secrets(text) - for secret in detected_secrets: - text = text.replace(secret["value"], "[REDACTED]") - if detected_secrets: - secret_types = [secret["type"] for secret in detected_secrets] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in message: {secret_types}" - ) - return text + inspected = walk_user_text(data, self.redact_text) + self._redact_prompt(data) - walk_user_text(data, _redact_message_text) + if inspected == 0: + # Image-only, empty-text, and unsupported payloads inspected + # nothing, so recording "allow" would count a run that never + # looked at any content. + _masked_entity_count.set(None) - if "prompt" in data: - if isinstance(data["prompt"], str): - detected_secrets = self.scan_message_for_secrets(data["prompt"]) - for secret in detected_secrets: - data["prompt"] = data["prompt"].replace( - secret["value"], "[REDACTED]" - ) - if len(detected_secrets) > 0: - secret_types = [secret["type"] for secret in detected_secrets] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in prompt: {secret_types}" - ) - elif isinstance(data["prompt"], list): - # Index back into the list — assigning to ``item`` would only - # rebind the loop variable and leave ``data["prompt"]`` - # carrying the unredacted secret. - for idx, item in enumerate(data["prompt"]): - if isinstance(item, str): - detected_secrets = self.scan_message_for_secrets(item) - for secret in detected_secrets: - item = item.replace(secret["value"], "[REDACTED]") - data["prompt"][idx] = item - if len(detected_secrets) > 0: - secret_types = [ - secret["type"] for secret in detected_secrets - ] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in prompt: {secret_types}" - ) - - # ``data["input"]`` (Responses API and embeddings/moderation) is - # already covered by ``walk_user_text`` above. return + + def _process_response( + self, + response: Optional[dict], + request_data: dict, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, + original_inputs: Optional[dict] = None, + ): + """Record allow/mask plus the masked-entity tally for a completed run. + + Records nothing when the guardrail inspected nothing (opted-out key, + empty inputs) or when the instance has no guardrail_name (legacy + ``litellm_settings.callbacks`` deployments, which predate guardrail + telemetry and stay without it). + """ + counts = _masked_entity_count.get() + _masked_entity_count.set(None) + if counts is None or self.guardrail_name is None: + return response + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="mask" if counts else "allow", + request_data=request_data, + guardrail_status="success", + duration=duration, + start_time=start_time, + end_time=end_time, + event_type=event_type, + guardrail_provider=GUARDRAIL_PROVIDER, + masked_entity_count=counts, + ) + return response + + def _process_error( + self, + e: Exception, + request_data: dict, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, + ): + """Label the failed run with this guardrail's provider so error rows + group with the successful ones in the monitor. Nameless legacy + instances record nothing, matching ``_process_response``.""" + _masked_entity_count.set(None) + if self.guardrail_name is None: + raise e + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=e, + request_data=request_data, + guardrail_status=( + "guardrail_intervened" + if self._is_guardrail_intervention(e) + else "guardrail_failed_to_respond" + ), + duration=duration, + start_time=start_time, + end_time=end_time, + event_type=event_type, + guardrail_provider=GUARDRAIL_PROVIDER, + ) + raise e 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/constants.py b/litellm/constants.py index c7b74e176db..ef9329b9dfc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b83e9b395a8..cb7da32f857 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -4,6 +4,7 @@ import logging import time from collections.abc import Mapping, Sequence from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response @@ -1164,6 +1165,12 @@ def _store_cost_breakdown_in_logging_obj( # Don't fail the main cost calculation if breakdown storage fails +def _without_provider_stated_cost(usage: Usage | None) -> Usage | None: + if usage is None or getattr(usage, "cost", None) is None: + return usage + return usage.model_copy(update=MappingProxyType({"cost": None})) + + def completion_cost( completion_response: object | None = None, model: str | None = None, @@ -1243,7 +1250,10 @@ def completion_cost( cache_creation_input_tokens: int | None = None cache_read_input_tokens: int | None = None audio_transcription_file_duration: float = 0.0 - cost_per_token_usage_object: Final[Usage | None] = _get_usage_object(completion_response=completion_response) + provider_usage_object: Final = _get_usage_object(completion_response=completion_response) + cost_per_token_usage_object: Final[Usage | None] = ( + _without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object + ) rerank_billed_units: RerankBilledUnits | None = None # Extract service_tier from optional_params if not provided directly 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/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 372c9bf6b91..c8e5610b4a9 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1485,7 +1485,7 @@ def log_guardrail_information(func): if func.__name__ == "apply_guardrail" and "inputs" in kwargs: original_inputs = kwargs.get("inputs") - logging_obj: Final = kwargs.get("logging_obj") + logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj") self_recorded_token: Final = _guardrail_self_recorded.set(False) try: response: Final = await func(*args, **kwargs) @@ -1527,7 +1527,7 @@ def log_guardrail_information(func): if func.__name__ == "apply_guardrail" and "inputs" in kwargs: original_inputs = kwargs.get("inputs") - logging_obj: Final = kwargs.get("logging_obj") + logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj") self_recorded_token: Final = _guardrail_self_recorded.set(False) try: response: Final = func(*args, **kwargs) 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/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2d737bc34e7..587da997f94 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,9 +10,9 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast -from typing_extensions import ReadOnly +from typing_extensions import Never, ReadOnly import litellm from litellm._logging import verbose_logger diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 480b1921c18..355faa41e68 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -54,6 +54,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) +_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) def _next_sync_or_exhausted(it: Any) -> object: @@ -1886,8 +1887,8 @@ class CustomStreamWrapper: @staticmethod def _resolve_provider_reported_cost(usage_cost: object) -> float | None: """ - Providers report usage.cost either as a number or, for Perplexity, as a - breakdown object whose total lives under ``total_cost``. + Providers report usage.cost either as a number or as a breakdown object + whose total lives under ``total_cost``. """ if isinstance(usage_cost, bool): return None @@ -1900,12 +1901,10 @@ class CustomStreamWrapper: @staticmethod def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", + custom_llm_provider: str | None, ) -> None: - """ - If the assembled response carries a provider-reported cost on - usage.cost, copy it into _hidden_params so litellm's cost - calculator uses it instead of a token-based estimate. - """ + if custom_llm_provider not in _USAGE_COST_HEADER_PROVIDERS: + return _usage: Final[Usage | None] = getattr(response, "usage", None) _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) if _cost is not None: @@ -2020,7 +2019,7 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) setattr( response, @@ -2270,7 +2269,7 @@ class CustomStreamWrapper: response: Final = self.model_response_creator() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) setattr( response, diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index f1c7451796d..5c30ff4747a 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation): return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + _, valid_parsed = self._parse_streaming_responses(responses_so_far) + combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed) + return StreamingScanKey(texts=(combined_text,)) + def _parse_streaming_responses( self, - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c7d12e5cf3a..c23797f72af 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( anthropic_tool_name, anthropic_tool_names, @@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @classmethod + def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(block) + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..b28daf73bc4 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -35,6 +35,22 @@ class StreamTransformSink: holdback_per_choice: dict[int, int] = field(default_factory=dict) +@dataclass(frozen=True, slots=True) +class StreamingScanKey: + """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys + compare equal when the round would scan the same content again; ``stream_ended`` + stays out of the comparison and only says whether the handler is on its + end-of-stream path, where an empty payload is still scanned today.""" + + texts: tuple[str, ...] + tool_calls: tuple[str, ...] = () + stream_ended: bool = field(default=False, compare=False) + + @property + def has_nothing_to_scan(self) -> bool: + return not self.stream_ended and not any(self.texts) and not self.tool_calls + + class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( @@ -151,6 +167,9 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + return None + def build_block_sse_chunks( self, exc: "ModifyResponseException", diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 9b6f9c47105..8dee262001d 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,6 +4,8 @@ import json from collections.abc import Callable, Iterator, Sequence from typing import Any, Final, TypeVar +from pydantic import BaseModel + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None: return getattr(item, field, None) +def stream_item_fingerprint(item: object) -> str: + plain: Final = item.model_dump() if isinstance(item, BaseModel) else item + return json.dumps(plain, sort_keys=True, default=str) + + +def stream_item_items(item: object, field: str) -> tuple[object, ...]: + value: Final = stream_item_field(item, field) + return tuple(value) if isinstance(value, (list, tuple)) else () + + def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: """ ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 8c306faa036..ac934ad0cb5 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -2,6 +2,7 @@ from typing import Final from httpx import Headers +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. - Deliberately does not fall back to `litellm_trace_id`: that is generated per - request (`str(uuid.uuid4())` when absent), so using it pins every request to a - different Fireworks node and prompt caching never hits. + Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the + proxy generated for a request that had none: both are per request, so using them + pins every request to a different Fireworks node and prompt caching never hits. """ params: Final = litellm_params + metadata: Final = params.get("metadata") + if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - metadata: Final = params.get("metadata") if isinstance(metadata, dict): value = metadata.get("session_id") if value: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 96a5ed663fc..d41c8557d72 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,6 +26,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( role_out_of_guardrail_scope, scoped_structured_message_indices, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here (see ``_process_streaming_transform`` for the incremental_diff path).""" - # check if the stream has ended - has_stream_ended = False - for chunk in responses_so_far: - if chunk.choices and chunk.choices[0].finish_reason is not None: - has_stream_ended = True - break + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: # convert to model response @@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) } + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + return StreamingScanKey( + texts=tuple(self._combine_streaming_texts(chunks).values()), + tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @staticmethod + def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(tool_call) + for chunk in responses_so_far + for choice in _stream_chunk_choices(chunk) + for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + ) + + @staticmethod + def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool: + first_choices: Final = tuple( + choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices + ) + return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices) + def _combine_streaming_texts( - self, responses_so_far: list["ModelResponseStream"] + self, responses_so_far: Sequence["ModelResponseStream"] ) -> dict[tuple[int, int | None], str]: """ Combine all streaming chunks into complete text per choice. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 5a5970fb867..a0db7aadb9e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation): ) return responses_so_far - def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + terminal_types: Final = frozenset( + ( + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + ) + ) + return stream_item_field(responses_so_far[-1], "type") in terminal_types + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + if not responses_so_far or not hasattr(responses_so_far[-1], "get"): + return None + last_event: Final = responses_so_far[-1] + last_event_type: Final = stream_item_field(last_event, "type") + if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: + return None + if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + return self._completed_response_scan_key(stream_item_field(last_event, "response")) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + stream_ended=self._check_streaming_has_ended(responses_so_far), + ) + + @staticmethod + def _completed_response_scan_key(response: object) -> StreamingScanKey: + output_items: Final = stream_item_items(response, "output") + message_items: Final = tuple( + item for item in output_items if stream_item_field(item, "type") != "function_call" + ) + return StreamingScanKey( + texts=tuple( + text + for item in message_items + for part in stream_item_items(item, "content") + if isinstance(text := stream_item_field(part, "text"), str) and text + ), + tool_calls=tuple( + stream_item_fingerprint(item) + for item in output_items + if stream_item_field(item, "type") == "function_call" + ), + stream_ended=True, + ) def build_stream_error_items( self, @@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation): ), ) - def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Get the string so far from the responses so far. @@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation): """ keyed_events: Final = tuple( ( - (event.get("item_id"), event.get("output_index"), event.get("content_index")), - event.get("text"), - event.get("delta"), + ( + stream_item_field(event, "item_id"), + stream_item_field(event, "output_index"), + stream_item_field(event, "content_index"), + ), + stream_item_field(event, "text"), + stream_item_field(event, "delta"), ) for event in responses_so_far - if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str) ) def part_text(part_key: tuple[object, object, object]) -> str: diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index ae5849812bf..09af8544b28 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -11,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, ) -from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.llms.xai.cost_calculator import ( apply_server_side_tool_usage_details_to_usage, ) @@ -30,6 +31,13 @@ from ...openai.chat.gpt_transformation import ( ) +def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None: + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if usage is None or reported_cost is None: + return None + return usage.model_copy(update=MappingProxyType({"cost": reported_cost})) + + class XAIChatConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: @@ -283,6 +291,9 @@ class XAIChatConfig(OpenAIGPTConfig): self._fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) + restated_usage: Final = _usage_restated_from_xai_ticks(getattr(response, "usage", None)) + if restated_usage is not None: + response.usage = restated_usage return response @staticmethod @@ -411,4 +422,8 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) - return super().chunk_parser(chunk) + parsed_chunk: Final = super().chunk_parser(chunk) + restated_usage: Final = _usage_restated_from_xai_ticks(getattr(parsed_chunk, "usage", None)) + if restated_usage is not None: + parsed_chunk.usage = restated_usage + return parsed_chunk diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index f122098332e..cf76a851a86 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -8,6 +8,17 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ProviderSpecificModelInfo +USD_TICKS_PER_DOLLAR: Final = 10_000_000_000 + + +def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None: + """xAI bills in ticks of a dollar: https://docs.x.ai/developers/cost-tracking""" + if not isinstance(cost_in_usd_ticks, int) or isinstance(cost_in_usd_ticks, bool): + return None + if cost_in_usd_ticks < 0: + return None + return cost_in_usd_ticks / USD_TICKS_PER_DOLLAR + class XAIModelInfo(BaseLLMModelInfo): def get_provider_info( diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index dd77b8d5d09..164568451d4 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -1,9 +1,11 @@ """ Helper util for handling XAI-specific cost calculation +- Prefers the cost xAI reports on the response over recomputing it locally - Uses the generic cost calculator which already handles tiered pricing correctly - Handles XAI-specific reasoning token billing (billed as part of completion tokens) """ +import math from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -36,6 +38,17 @@ def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping usage.prompt_tokens_details = prompt_tokens_details # rebind-ok: write details onto caller usage +def _cost_reported_by_xai(usage: "Usage") -> float | None: + reported_cost: Final[object] = getattr(usage, "cost", None) + if not isinstance(reported_cost, (int, float)) or isinstance(reported_cost, bool): + return None + if not math.isfinite(reported_cost): + return None + if reported_cost < 0: + return None + return float(reported_cost) + + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. @@ -48,6 +61,10 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ + reported_cost: Final = _cost_reported_by_xai(usage) + if reported_cost is not None: + return 0.0, reported_cost + # XAI-specific completion cost: completion is billed as visible + reasoning # tokens. Detect when the transformation layer already folded them so we # don't double-count; fall back to raw xAI shape for callers that bypass @@ -112,6 +129,9 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa Per-call rate comes from model_info.search_context_cost_per_query when set, otherwise the default xAI tools rate ($5 / 1k calls). """ + if _cost_reported_by_xai(usage) is not None: + return 0.0 + details: Final = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): return 0.0 diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index d79e7d4c146..acf03d88911 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,17 +1,44 @@ -from typing import Any, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final + +import httpx import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig -from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as _LiteLLMLoggingObj, + ) + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +def _usage_restated_from_xai_ticks(usage: ResponseAPIUsage | None) -> ResponseAPIUsage | None: + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if usage is None or reported_cost is None: + return None + return usage.model_copy(update=MappingProxyType({"cost": reported_cost})) + class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ @@ -250,6 +277,41 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return f"{api_base}/responses" + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + response: Final = super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + restated_usage: Final = _usage_restated_from_xai_ticks(response.usage) + if restated_usage is not None: + response.usage = restated_usage + return response + + def transform_streaming_response( + self, + model: str, + parsed_chunk: dict, # mutable-ok: overrides the base class signature + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIStreamingResponse: + event: Final = super().transform_streaming_response( + model=model, + parsed_chunk=parsed_chunk, + logging_obj=logging_obj, + ) + if not isinstance(event, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent)): + return event + restated_usage: Final = _usage_restated_from_xai_ticks(event.response.usage) + if restated_usage is not None: + event.response.usage = restated_usage + return event + def supports_native_websocket(self) -> bool: """XAI does not support native WebSocket for Responses API""" return False diff --git a/litellm/main.py b/litellm/main.py index 0128e4defe5..0bca4a7350e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8595,9 +8595,19 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N return TextCompletionResponse(**response) +_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) + + +def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: + if logging_obj is None: + return False + provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider") + return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS + + def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None: usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) - if isinstance(usage_cost, (int, float)): + if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj): return float(usage_cost) if logging_obj is not None: return None 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/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 4f6305d88cf..af02c11ad86 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -5,10 +5,10 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never +from typing import TYPE_CHECKING, Any, Final, TypedDict from pydantic import ValidationError -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2da7ceb2d50..849e54c65aa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", ) + missing_session_id: Literal["generate", "reject"] | None = Field( + None, + description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", + ) enable_public_model_hub: bool = Field( default=False, description="Public model hub for users to see what models they have access to, supported openai params, etc.", diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 3e4dc07a521..3b8151d1064 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -13,10 +13,10 @@ import os import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Annotated, Final, TypedDict, assert_never +from typing import Annotated, Final, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1682cf12f4e..47f69732e95 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,7 +7,9 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar, assert_never +from typing import Final, Literal, Protocol, TypeVar + +from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger 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_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 3d2ed641a30..744b2959c73 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -8,7 +8,7 @@ import json import os from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone -from types import UnionType +from types import MappingProxyType, UnionType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse @@ -51,6 +51,9 @@ from litellm.types.guardrails import ( SupportedGuardrailIntegrations, ToolPermissionGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.hide_secrets import ( + HideSecretsGuardrailConfigModel, +) if TYPE_CHECKING: from types import CodeType @@ -1401,7 +1404,11 @@ async def get_guardrail_ui_settings(): provider: [hook.value for hook in hooks] for provider, guardrail_class in guardrail_class_registry.items() if (hooks := guardrail_class.get_supported_event_hooks()) is not None - } + } | MappingProxyType( + # hide-secrets lives in the enterprise package, not in the registry + # above; it only runs on pre_call. + {SupportedGuardrailIntegrations.HIDE_SECRETS.value: [GuardrailEventHooks.pre_call.value]} + ) return GuardrailUIAddGuardrailSettings( supported_entities=[entity.value for entity in PiiEntityType], @@ -1953,12 +1960,18 @@ async def get_provider_specific_params(): tool_permission_fields["ui_friendly_name"] = ToolPermissionGuardrailConfigModel.ui_friendly_name() + # hide-secrets lives in the enterprise package, not in the registry loop below. + hide_secrets_fields: Final = _get_fields_from_model(HideSecretsGuardrailConfigModel) + + hide_secrets_fields["ui_friendly_name"] = HideSecretsGuardrailConfigModel.ui_friendly_name() + # Return the provider-specific parameters provider_params: Final = { SupportedGuardrailIntegrations.BEDROCK.value: bedrock_fields, SupportedGuardrailIntegrations.PRESIDIO.value: presidio_fields, SupportedGuardrailIntegrations.LAKERA_V2.value: lakera_v2_fields, SupportedGuardrailIntegrations.TOOL_PERMISSION.value: tool_permission_fields, + SupportedGuardrailIntegrations.HIDE_SECRETS.value: hide_secrets_fields, } ### get the config model for the guardrail - go through the registry and get the config model for the guardrail 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/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 46b00829b74..c6b8df1b493 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, ) # Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error @@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol): @property def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + @property + def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... + @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... @@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]: return choices +def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool: + if scan_key is None: + return False + return scan_key == last_scan_key or scan_key.has_nothing_to_scan + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" @@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False + last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round async for item in response: chunk_counter += 1 @@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", + chunk_counter, + guardrail_to_apply.guardrail_name, + ) + chunks_yielded = True + responses_yielded.append(item) + yield item + continue + verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, @@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger): # string, permanently losing this chunk's content. original_item = copy.deepcopy(item) - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, @@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger): ): yield error_item return + if scan_key is not None: + last_scan_key = scan_key chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger): # preserve the list, not clone every chunk (deepcopy would double # peak memory for large responses). buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(end_scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", + guardrail_to_apply.guardrail_name, + ) + for buffered_item in buffered_items or (): + yield buffered_item + for pending_item in pending_end_of_stream_items: + responses_yielded.append(pending_item) + yield pending_item + return try: await endpoint_translation.process_output_streaming_response( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 20f83085286..1d440448c2f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -16,6 +16,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm._uuid import uuid from litellm.constants import ( CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -23,6 +24,7 @@ from litellm.constants import ( OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -40,6 +42,7 @@ from litellm.proxy._types import ( AddTeamCallback, CommonProxyErrors, LitellmDataForBackendLLMCall, + LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -47,6 +50,8 @@ from litellm.proxy._types import ( TeamCallbackMetadata, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import get_request_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, @@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: return session_id +def _is_llm_inference_route(request: Request) -> bool: + route: Final = get_request_route(request) + return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) + + +def apply_missing_session_id_policy( + data: dict[str, object], # mutable-ok: stamps session ids in place on the request body the pipeline threads through + _metadata_variable_name: str, + general_settings: Mapping[str, object] | None, + request: Request, +) -> None: + policy: Final = general_settings.get("missing_session_id") if general_settings else None + if policy is None or not _is_llm_inference_route(request): + return + metadata: Final = data.get(_metadata_variable_name) + if not isinstance(metadata, dict): + return + if data.get("litellm_session_id") or metadata.get("session_id"): + return + match policy: + case "generate": + session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4()) + data["litellm_session_id"] = session_id # rebind-ok: data is an out-param + data.setdefault("litellm_trace_id", session_id) + metadata["session_id"] = session_id + metadata[SESSION_ID_GENERATED_METADATA_KEY] = True + case "reject": + raise ProxyException( + message=( + "Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. " + "Required by `general_settings.missing_session_id: reject`." + ), + type=ProxyErrorTypes.bad_request_error, + param="session_id", + code=400, + ) + case _: + verbose_proxy_logger.warning( + "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy + ) + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code identifies itself as ``claude-cli/ ...``; the IDE extensions and the Agent SDK run through the same CLI and share that prefix.""" @@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request( data=data, _metadata_variable_name=_metadata_variable_name, ) + apply_missing_session_id_policy( + data=data, + _metadata_variable_name=_metadata_variable_name, + general_settings=general_settings, + request=request, + ) # Expose request headers under the metadata field for guardrails (fixes #17477) if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict): diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 466b100ea1f..12e1f1a03f3 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -443,7 +443,9 @@ class SAMLAuthHandler: last_name: Final = SAMLAuthHandler._attribute_value( attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES ) - role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES) + role_values: Final = SAMLAuthHandler._attribute_values( + attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES + ) team_ids: Final = SAMLAuthHandler._attribute_values( attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES ) @@ -464,7 +466,7 @@ class SAMLAuthHandler: picture=None, provider="saml", team_ids=team_ids, - user_role=get_litellm_user_role(role_value) if role_value else None, + user_role=get_litellm_user_role(role_values), ) except ValidationError as e: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 070df97d09a..4414eed97b2 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -4,12 +4,44 @@ Types for the management endpoints Might include fastapi/proxy requirements.txt related imports """ +from collections.abc import Iterable, Sequence from typing import Any, Final, cast from fastapi_sso.sso.base import OpenID from litellm.proxy._types import LitellmUserRoles +# Ordered highest to lowest privilege +LITELLM_USER_ROLE_HIERARCHY: Final = ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, +) + + +def highest_privilege_role(roles: Iterable[LitellmUserRoles]) -> LitellmUserRoles | None: + """ + Pick the highest privilege role out of the roles an IdP asserted for one user. + + IdPs do not guarantee ordering within a multi-valued role claim, so a user holding + several roles resolves to the most privileged one rather than whichever came first. + Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay + deterministic. + + Args: + roles: The roles resolved from the claim + + Returns: + The highest privilege role, or None if `roles` is empty + """ + resolved: Final = frozenset(roles) + if not resolved: + return None + + ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) + return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + def is_valid_litellm_user_role(role_str: str) -> bool: """ @@ -28,12 +60,22 @@ def is_valid_litellm_user_role(role_str: str) -> bool: return False -def get_litellm_user_role(role_str) -> LitellmUserRoles | None: +def _role_from_claim_value(role_str: object) -> LitellmUserRoles | None: + if not isinstance(role_str, str): + return None + # Use _value2member_map_ for O(1) lookup, case-insensitive + result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) + return cast(LitellmUserRoles | None, result) + + +def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: """ Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive). Handles list inputs since some SSO providers (e.g., Keycloak) return roles - as arrays like ["proxy_admin"] instead of plain strings. + as arrays like ["proxy_admin"] instead of plain strings. A claim carrying several + roles resolves to the highest privilege one, so a user does not lose access just + because the IdP listed a weaker role first. Args: role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"]) @@ -41,16 +83,12 @@ def get_litellm_user_role(role_str) -> LitellmUserRoles | None: Returns: LitellmUserRoles enum if valid, None otherwise """ - try: - if isinstance(role_str, list): - if len(role_str) == 0: - return None - role_str = role_str[0] - # Use _value2member_map_ for O(1) lookup, case-insensitive - result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) - return cast(LitellmUserRoles | None, result) - except Exception: - return None + if isinstance(role_str, (list, tuple)): + entries: Final = cast(Sequence[object], role_str) # cast-ok: isinstance narrows the claim, not its elements + return highest_privilege_role( + role for role in (_role_from_claim_value(entry) for entry in entries) if role is not None + ) + return _role_from_claim_value(role_str) class CustomOpenID(OpenID): diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1feefa5725d..84150ef7935 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.sso_helper_utils import ( ) from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add from litellm.proxy.management_endpoints.types import ( + LITELLM_USER_ROLE_HIERARCHY, CustomOpenID, get_litellm_user_role, is_valid_litellm_user_role, @@ -809,15 +810,6 @@ def normalize_email(email: str | None) -> str | None: return email.lower() if isinstance(email, str) else email -# Ordered highest to lowest privilege -LITELLM_USER_ROLE_HIERARCHY: Final = ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, -) - - def determine_role_from_groups( user_groups: list[str], role_mappings: "RoleMappings", @@ -4312,14 +4304,7 @@ class MicrosoftSSOHandler: listed first. Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay deterministic """ - resolved: Final = frozenset( - role for role in (get_litellm_user_role(role_str) for role_str in app_roles or ()) if role is not None - ) - if not resolved: - return None - - ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) - return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + return get_litellm_user_role(tuple(app_roles or ())) @staticmethod def get_app_roles_from_id_token(id_token: str | None) -> list[str]: diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 0aee5e8cc54..a41bd36d510 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -2,7 +2,9 @@ import json from collections.abc import Iterator from dataclasses import dataclass from itertools import chain -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 9d450cb5b8d..8c59a520272 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose. from dataclasses import dataclass from pathlib import Path -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.path_utils import safe_filename diff --git a/litellm/proxy/public_endpoints/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json index 1e0d40567d7..6a6642d4211 100644 --- a/litellm/proxy/public_endpoints/autorouter_presets.json +++ b/litellm/proxy/public_endpoints/autorouter_presets.json @@ -1,4 +1,26 @@ { + "1m_context": { + "label": "1M Context", + "description": "Routes across models with 1M-token context windows: Luna for simple queries, Terra for medium, Opus 5 for complex, Opus 5 at high thinking for reasoning.", + "complexity_router_config": { + "tiers": { + "SIMPLE": ["gpt-5.6-luna"], + "MEDIUM": ["gpt-5.6-terra"], + "COMPLEX": ["claude-opus-5"], + "REASONING": ["claude-opus-5"] + }, + "tier_model_configs": { + "REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }] + }, + "classifier_type": "heuristic_v2", + "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", + "session_affinity": false, + "modality_routing": false, + "modality_pin_override": false, + "deployment_affinity": true + } + }, "anthropic_family": { "label": "Anthropic Family", "description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Opus at high thinking for reasoning.", @@ -17,6 +39,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } }, @@ -35,6 +58,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } }, @@ -63,6 +87,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } }, @@ -84,6 +109,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 d205db90607..a00ae6bee80 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, create_model from litellm._logging import verbose_router_logger -from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import ( + EMPTY_MAPPING, + RETURN_RAW_MODEL_NAME_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata @@ -747,7 +751,8 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo A modality escalation is transient the same way: it describes what this one call carries (an image), not what the session's traffic looks like, and pinning it would hold every following - text turn on the vision-capable model the image forced. + text turn on the vision-capable model the image forced. A modality pin override is the same + fact on a session that already holds a pin, so it must not overwrite the pin it displaced. """ return decision is None or ( decision.get("cause") @@ -756,6 +761,7 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo "plan_mode", "housekeeping", "modality_escalation", + "modality_pin_override", ) and not decision.get("context_escalated") ) @@ -2389,8 +2395,11 @@ class ComplexityRouter(CustomLogger): """Replace a routed model that cannot accept this request's image input. The single modality owner, applied to the decided response at the hook's exits so every - routing path is covered uniformly. A KEPT session pin is exempt by design (its cause); - replacement picks and every other path are just responses. The re-placement walks + routing path is covered uniformly. A KEPT session pin is exempt by design (its cause) + unless modality_pin_override is set, in which case the image turn is re-placed and reported + as modality_pin_override while the stored pin, written upstream from the session's own + model, is left for the next text turn; replacement picks and every other path are just + responses. The re-placement walks UPWARD-ONLY from the decision's tier (so a plan-mode floor can never be undercut), picks through `_pick_model_for_tier` so routing plugins still apply, then falls to default_model (never on plugin routers, and never on a plan-floored decision, since @@ -2403,7 +2412,11 @@ class ComplexityRouter(CustomLogger): not self.config.modality_routing or not resolved_messages or response.model is None - or (decision is not None and decision.get("cause") == "session_affinity_pin") + or ( + decision is not None + and decision.get("cause") == "session_affinity_pin" + and not self.config.modality_pin_override + ) or not request_contains_image_content(resolved_messages) or self._model_accepts_image_input(response.model) ): @@ -2445,6 +2458,10 @@ class ComplexityRouter(CustomLogger): self._restamp_adaptive_choice(request_kwargs, response.model, new_model) same_tier: Final = capable is not None and decided == capable base_cause: Final = (decision.get("cause") if decision is not None else None) or "default_fallback" + # Reaching here on a kept pin means modality_pin_override is on, since the guard above + # returns otherwise. The model moved off the pin even on a same-tier repick, so reporting + # the pin's own cause would claim the session's model served a request it did not. + displaced_pin: Final = base_cause == "session_affinity_pin" displaced_default: Final = decided is None and response.model == self.config.default_model markers: Final = ( "modality:image", @@ -2454,7 +2471,7 @@ class ComplexityRouter(CustomLogger): old_signals: Final = tuple(decision.get("signals") or ()) if decision is not None else () new_decision: Final = self._build_routing_decision( routed_model=new_model, - cause=base_cause if same_tier else "modality_escalation", + cause="modality_pin_override" if displaced_pin else (base_cause if same_tier else "modality_escalation"), tier=new_tier, score=decision.get("score") if decision is not None else None, signals=(*old_signals, *markers), @@ -2712,7 +2729,7 @@ class ComplexityRouter(CustomLogger): """Resolve a client-supplied session_id.""" for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): session_id = metadata.get("session_id") - if session_id is not None: + if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return str(session_id) return None 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/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b1e9dbdefa8..39d3e25aacb 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -21,7 +21,7 @@ from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") - if session_id is None: + if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None return str(session_id) 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/llms/gemini_audio_transcription.py b/litellm/types/llms/gemini_audio_transcription.py index cb12e0f45b8..f7e74ba4bf8 100644 --- a/litellm/types/llms/gemini_audio_transcription.py +++ b/litellm/types/llms/gemini_audio_transcription.py @@ -1,7 +1,7 @@ -from typing import Literal, Required +from typing import Literal from pydantic import BaseModel, ConfigDict -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict class GeminiTranscriptionAudioInput(TypedDict): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py b/litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py new file mode 100644 index 00000000000..3100968e8b1 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py @@ -0,0 +1,20 @@ +"""Types for the Hide Secrets guardrail.""" + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class HideSecretsGuardrailConfigModel(GuardrailConfigModel): + """Configuration for the Hide Secrets guardrail. Detection runs in-process + on the detect-secrets library; ``detect_secrets_config`` overrides the + bundled plugin set.""" + + detect_secrets_config: dict | None = Field( # mutable-ok: UI type derivation maps dict to "object" + default=None, + description="Optional detect-secrets configuration (plugins_used, filters_used) overriding the bundled plugin set", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Hide Secrets" 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/ruff-strict-budget.json b/ruff-strict-budget.json index 4fcf650a8bc..be2b30fc189 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2985 + "limit": 2984 }, "ANN002": { "limit": 71 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2917 + "limit": 2916 }, "C401": { "limit": 8 diff --git a/tests/code_coverage_tests/check_py310_typing_imports.py b/tests/code_coverage_tests/check_py310_typing_imports.py new file mode 100644 index 00000000000..0cd4d089890 --- /dev/null +++ b/tests/code_coverage_tests/check_py310_typing_imports.py @@ -0,0 +1,150 @@ +import ast +import os +import sys +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset( + { + "NotRequired", + "Required", + "Self", + "LiteralString", + "Never", + "assert_never", + "assert_type", + "reveal_type", + "TypeVarTuple", + "Unpack", + "dataclass_transform", + "override", + "TypeAliasType", + "get_original_bases", + "ReadOnly", + "TypeIs", + "NoDefault", + "get_protocol_members", + "is_protocol", + "evaluate_forward_ref", + "TypeForm", + } +) + + +@dataclass(frozen=True, slots=True) +class TypingImportViolation: + file: str + line: int + name: str + + +def _walk_with_ancestors( + node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = () +) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]: + yield node, ancestors + for field_name, field_value in ast.iter_fields(node): + if isinstance(field_value, ast.AST): + yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name))) + elif isinstance(field_value, list): + for child in field_value: + if isinstance(child, ast.AST): + yield from _walk_with_ancestors(child, (*ancestors, (node, field_name))) + + +def _is_sys_version_info(node: ast.AST) -> bool: + return ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "sys" + and node.attr == "version_info" + ) + + +def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool: + nearest_if: Final[tuple[ast.If, str] | None] = next( + ( + (ancestor, field_name) + for ancestor, field_name in reversed(ancestors) + if isinstance(ancestor, ast.If) + ), + None, + ) + if nearest_if is None: + return False + enclosing_if, branch = nearest_if + test: Final[ast.expr] = enclosing_if.test + if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left): + return False + operator: Final[ast.cmpop] = test.ops[0] + return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or ( + isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse" + ) + + +def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]: + path: Final[Path] = Path(file_path) + tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + return tuple( + violation + for node, ancestors in _walk_with_ancestors(tree) + if not _is_version_guarded(ancestors) + for violation in _violations_for_node(node, path) + ) + + +def _violations_for_node( + node: ast.AST, path: Path +) -> tuple[TypingImportViolation, ...]: + if isinstance(node, ast.ImportFrom) and node.module == "typing": + return tuple( + TypingImportViolation(file=str(path), line=node.lineno, name=alias.name) + for alias in node.names + if alias.name in PY311_PLUS_TYPING_NAMES + ) + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "typing" + and node.attr in PY311_PLUS_TYPING_NAMES + ): + return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),) + return () + + +def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]: + base_path: Final[Path] = Path(base_dir) + return tuple( + violation + for directory in ( + base_path / "litellm", + base_path / "enterprise", + base_path / "litellm-proxy-extras" / "litellm_proxy_extras", + ) + if directory.exists() + for path in directory.rglob("*.py") + for violation in scan_file(path) + ) + + +def main() -> None: + violations: Final[tuple[TypingImportViolation, ...]] = scan_directory() + if violations: + message: Final[str] = "\n".join( + ( + "Python 3.10-incompatible typing imports found:", + *( + f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; " + "import it from typing_extensions instead because litellm supports Python 3.10" + for violation in violations + ), + ) + ) + sys.stdout.write(f"{message}\n") + raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead") + sys.stdout.write("No Python 3.10-incompatible typing imports found.\n") + + +if __name__ == "__main__": + main() diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 044e1e0de8e..9faaaf492e8 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -370,12 +370,12 @@ 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, client=None, _is_async=False, - embedding_executor=None, ): litellm_params_dict = ( litellm_params.model_dump(exclude_none=False) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 1a7fb1f3e41..abae26e02cd 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,5 +1,5 @@ import httpx -from openai import OpenAI, BadRequestError, APIStatusError +from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError import pytest @@ -105,10 +105,9 @@ def test_streaming_response(): assert len(collected_chunks) > 0 -def test_bad_request_error(): +def test_model_not_found_error(): client = get_test_client() - with pytest.raises(BadRequestError): - # Trigger error with invalid model name + with pytest.raises(NotFoundError): client.responses.create(model="non-existent-model", input="This should fail") 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/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py new file mode 100644 index 00000000000..dc1cbb9983e --- /dev/null +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -0,0 +1,274 @@ +"""Tests for the hide-secrets guardrail (LIT-3548). + +Covers the three defects from the ticket: +- ``apply_guardrail`` (the UI test playground path) must redact, not echo. +- Guardrail runs must record ``standard_logging_guardrail_information`` so + Spend Logs / the guardrails monitor show activity, with hits ("mask" + + masked_entity_count) distinguishable from clean requests ("allow"). +- Defining ``apply_guardrail`` must NOT reroute proxied traffic off the + native ``async_pre_call_hook`` (per-key opt-out and ``data["prompt"]`` + handling live only on the native path). +""" + +import pytest + +from litellm_enterprise.enterprise_callbacks.secret_detection import ( + _ENTERPRISE_SecretDetection, +) +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth + +AWS_KEY = "AKIAIOSFODNN7EXAMPLE" + + +def _guardrail() -> _ENTERPRISE_SecretDetection: + return _ENTERPRISE_SecretDetection( + guardrail_name="hide-secrets", event_hook="pre_call", default_on=True + ) + + +def _recorded(request_data: dict) -> dict: + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(entries) == 1 + return entries[0] + + +@pytest.mark.asyncio +async def test_apply_guardrail_redacts_secrets(): + """Playground path: the returned texts must carry [REDACTED], not the secret.""" + guardrail = _guardrail() + request_data: dict = {"metadata": {}} + + result = await guardrail.apply_guardrail( + inputs={"texts": [f"my key is {AWS_KEY}, keep it safe"]}, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == ["my key is [REDACTED], keep it safe"] + + recorded = _recorded(request_data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "mask" + assert recorded["guardrail_provider"] == "hide-secrets" + assert recorded["masked_entity_count"] == {"AWS Access Key": 1} + + +@pytest.mark.asyncio +async def test_apply_guardrail_clean_text_records_allow(): + guardrail = _guardrail() + request_data: dict = {"metadata": {}} + + result = await guardrail.apply_guardrail( + inputs={"texts": ["nothing sensitive here"]}, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == ["nothing sensitive here"] + + recorded = _recorded(request_data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "allow" + assert recorded["masked_entity_count"] == {} + + +@pytest.mark.asyncio +async def test_pre_call_hook_records_mask_with_entity_count(): + """Live-traffic path: a redaction must be visible in spend-log telemetry.""" + guardrail = _guardrail() + data = { + "messages": [{"role": "user", "content": f"use {AWS_KEY} for auth"}], + "metadata": {}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["messages"][0]["content"] == "use [REDACTED] for auth" + + recorded = _recorded(data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "mask" + assert recorded["guardrail_provider"] == "hide-secrets" + assert recorded["masked_entity_count"] == {"AWS Access Key": 1} + + +@pytest.mark.asyncio +async def test_pre_call_hook_clean_request_records_allow(): + """A request with no secrets must be distinguishable from a redacted one.""" + guardrail = _guardrail() + data = { + "messages": [{"role": "user", "content": "what's the weather"}], + "metadata": {}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + recorded = _recorded(data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "allow" + assert recorded["masked_entity_count"] == {} + + +@pytest.mark.asyncio +async def test_pre_call_hook_opt_out_records_nothing(): + """A key with permissions={"hide_secrets": False} skips redaction, so no + telemetry is recorded: every reader of a recorded entry (guardrail usage + tracking, compliance checks, the spend-log viewer) counts it as a run.""" + guardrail = _guardrail() + content = f"my key is {AWS_KEY}" + data = {"messages": [{"role": "user", "content": content}], "metadata": {}} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(permissions={"hide_secrets": False}), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["messages"][0]["content"] == content # untouched + assert "standard_logging_guardrail_information" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_still_redacts_text_completion_prompt(): + """data["prompt"] (str and list) is a native-hook-only surface; it must + keep redacting now that the class also implements apply_guardrail.""" + guardrail = _guardrail() + data = {"prompt": f"key {AWS_KEY} end", "metadata": {}} + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + assert data["prompt"] == "key [REDACTED] end" + + guardrail = _guardrail() + data = {"prompt": [f"key {AWS_KEY}", "clean"], "metadata": {}} + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + assert data["prompt"] == ["key [REDACTED]", "clean"] + + +def test_proxied_traffic_stays_on_native_hooks(): + """Implementing apply_guardrail must not reroute proxied requests onto the + unified path: that path skips ``should_run_check`` (per-key opt-out) and + never sees ``data["prompt"]``.""" + guardrail = _guardrail() + assert guardrail.uses_apply_guardrail_interface() is True + assert guardrail._deployment_pre_call_target() is guardrail + + +@pytest.mark.asyncio +async def test_apply_guardrail_without_texts_records_nothing(): + """No inputs means nothing was inspected, so no "allow" row is recorded. + Empty strings count as no input: there is no content to inspect.""" + guardrail = _guardrail() + + empty_variants: list[list[str]] = [[], ["", ""]] + for texts in empty_variants: + request_data: dict = {"metadata": {}} + result = await guardrail.apply_guardrail( + inputs={"texts": texts}, request_data=request_data, input_type="request" + ) + assert result == {"texts": texts} + assert "standard_logging_guardrail_information" not in request_data["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data", + [ + pytest.param( + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://x/y.png"}} + ], + } + ], + "metadata": {}, + }, + id="image_only", + ), + pytest.param( + {"messages": [{"role": "user", "content": ""}], "metadata": {}}, + id="empty_message", + ), + pytest.param({"prompt": "", "metadata": {}}, id="empty_prompt"), + pytest.param({"prompt": ["", ""], "metadata": {}}, id="empty_prompt_list"), + ], +) +async def test_pre_call_hook_without_inspectable_text_records_nothing(data: dict): + """A payload the guardrail could not inspect (image-only content, empty + strings) must not record an "allow" run: monitoring would count a check + that never looked at any text.""" + guardrail = _guardrail() + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert "standard_logging_guardrail_information" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_mixed_prompt_list_still_redacts_and_records(): + """A prompt list mixing empty and real strings is inspected, so the run is + recorded and the non-empty entry is still redacted.""" + guardrail = _guardrail() + data = {"prompt": ["", f"key {AWS_KEY}"], "metadata": {}} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["prompt"] == ["", "key [REDACTED]"] + recorded = _recorded(data) + assert recorded["guardrail_response"] == "mask" + assert recorded["masked_entity_count"] == {"AWS Access Key": 1} + + +@pytest.mark.asyncio +async def test_legacy_nameless_instance_records_nothing(): + """``litellm_settings.callbacks: ["hide_secrets"]`` builds an arg-less + instance with no guardrail_name. It still redacts, but recording a nameless + entry would flip every spend row's guardrail status with nothing to join on.""" + guardrail = _ENTERPRISE_SecretDetection() + data = { + "messages": [{"role": "user", "content": f"use {AWS_KEY} for auth"}], + "metadata": {}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["messages"][0]["content"] == "use [REDACTED] for auth" + assert "standard_logging_guardrail_information" not in data["metadata"] 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/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 4c99bce2b0f..4b2a3f9247c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.streaming_handler import ( from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, + ModelResponse, ModelResponseStream, PromptTokensDetailsWrapper, StandardLoggingPayload, @@ -1750,7 +1751,7 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert complete_response.usage.cost == 0.00025 # Use the real propagation method from CustomStreamWrapper - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openrouter") assert "additional_headers" in complete_response._hidden_params assert ( @@ -1769,14 +1770,12 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert provider_cost == 0.00025 -def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): - """ - Regression: Perplexity reports usage.cost as a breakdown object, which used to - blow up the end of the stream with - `float() argument must be a string or a real number, not 'dict'`. - """ +def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator(): import litellm - from litellm.cost_calculator import get_response_cost_from_hidden_params + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, + ) chunks = [ ModelResponseStream( @@ -1828,13 +1827,81 @@ def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): assert complete_response is not None - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "perplexity") - assert ( - get_response_cost_from_hidden_params(complete_response._hidden_params) - == 0.00503 + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + assert response_cost_calculator( + response_object=complete_response, + model="perplexity/sonar", + custom_llm_provider="perplexity", + call_type="completion", + optional_params={}, + ) == pytest.approx(0.00503) + + +def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map(): + import litellm + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, ) + model = "openai/streams-cost-in-nanodollars" + litellm.register_model( + { + model: { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + complete_response = ModelResponse( + id="chatcmpl-openai-compatible", + model=model, + choices=[], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000), + ) + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai") + + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + assert response_cost_calculator( + response_object=complete_response, + model=model, + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) == pytest.approx(2e-5) + + +def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch): + import litellm + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, + ) + + complete_response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[], + usage=Usage(completion_tokens=353, prompt_tokens=198, total_tokens=551, cost=0.0009956), + ) + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "xai") + + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + assert response_cost_calculator( + response_object=complete_response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + call_type="completion", + optional_params={}, + ) == pytest.approx(0.0009956 * 1.5) + def test_provider_reported_cost_ignores_unusable_shapes(): assert CustomStreamWrapper._resolve_provider_reported_cost(None) is None diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py new file mode 100644 index 00000000000..dd7e8fadcd8 --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -0,0 +1,35 @@ +"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" + +import json + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + +def _text_event(text: str) -> str: + return json.dumps( + { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]}, + } + ) + + +def _status_event() -> str: + return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}}) + + +class TestA2AGuardrailHandlerStreamingScanKey: + def test_key_joins_the_text_of_every_message_event(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")]) + assert key == StreamingScanKey(texts=("hello world",)) + + def test_events_without_text_leave_the_key_unchanged(self): + handler = A2AGuardrailHandler() + events = [_text_event("hello")] + assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events) + + def test_unparseable_items_are_ignored(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 0fe7730e91e..3044a321aa6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) @@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults: } later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)] assert {"type": "text", "text": "Now fetch the page."} in later_blocks + + +class TestAnthropicMessagesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _sse(event_type, data): + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + def _text_delta(self, text): + return self._sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + def test_key_is_empty_before_any_text_arrives(self): + head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}}) + key = AnthropicMessagesHandler().get_streaming_scan_key([head]) + assert key == StreamingScanKey(texts=("",)) + + def test_key_accumulates_text_deltas(self): + key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")]) + assert key.texts == ("hello world",) + assert key.stream_ended is False + + def _stop(self, stop_reason): + return self._sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}}, + ) + + def test_stop_without_tool_use_scans_the_same_payload(self): + handler = AnthropicMessagesHandler() + open_key = handler.get_streaming_scan_key([self._text_delta("hi")]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self): + handler = AnthropicMessagesHandler() + tool_use = self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}, + }, + ) + open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) + assert open_key == StreamingScanKey(texts=("hi",)) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key 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/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d7cc89868af..ec8725db5f7 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -8,6 +8,7 @@ import litellm from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): ) +def test_get_fireworks_session_id_ignores_proxy_generated_session_id(): + """general_settings.missing_session_id: generate stamps a fresh id per request; sending it + as x-session-affinity would pin every request to a different node.""" + assert ( + get_fireworks_session_id( + { + "litellm_session_id": "generated-1", + "litellm_trace_id": "generated-1", + "metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + } + ) + is None + ) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 7dd6065063a..cebab2512d0 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -12,6 +12,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) @@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded: ) ] assert handler._check_streaming_has_ended(chunks) is True + + +class TestStreamingScanKey: + """get_streaming_scan_key identifies what a sampled round would scan so the + unified hook can skip rounds that would re-scan already-cleared text""" + + @staticmethod + def _chunk(content, finish_reason=None, index=0): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + def test_key_carries_accumulated_text_and_open_stream(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_chunks_without_text_leave_the_key_unchanged(self): + handler = OpenAIChatCompletionsHandler() + before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)]) + assert after == before + + def test_finish_chunk_without_tool_calls_scans_the_same_payload(self): + handler = OpenAIChatCompletionsHandler() + open_key = handler.get_streaming_scan_key([self._chunk("hi")]) + ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")]) + assert open_key.stream_ended is False + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self): + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + handler = OpenAIChatCompletionsHandler() + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + tool_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] + ) + assert open_key == StreamingScanKey(texts=("hi",)) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): + handler = OpenAIChatCompletionsHandler() + first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] + key_at_first_finish = handler.get_streaming_scan_key(first_done) + key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)]) + assert key_at_first_finish.stream_ended is True + assert key_after_more_text.stream_ended is True + assert key_after_more_text != key_at_first_finish + + def test_non_stream_items_are_ignored(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index d071ef78c2d..295121167d6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks: dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] assert len(dones) == 1 assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." + + +class TestOpenAIResponsesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + def test_no_events_yields_no_key(self): + assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None + + def test_key_accumulates_deltas_while_the_stream_is_open(self): + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_typed_delta_events_accumulate_like_dicts(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + events = [ + OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=i, + ) + for i, text in enumerate(("hel", "lo")) + ] + key = OpenAIResponsesHandler().get_streaming_scan_key(events) + assert key.texts == ("hello",) + assert key.stream_ended is False + + def test_events_without_text_leave_the_key_unchanged(self): + handler = OpenAIResponsesHandler() + events = [self._delta(0, "hi")] + quiet = events + [{"type": "response.in_progress", "sequence_number": 1}] + assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events) + + @staticmethod + def _completed(sequence_number, output): + return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}} + + def test_completed_event_keys_on_the_final_output_text(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_completed_event_with_a_function_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])]) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_completed_event_reads_every_output_text_part(self): + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + item = GenericResponseOutputItem( + type="message", + id="msg_1", + status="completed", + role="assistant", + content=[ + OutputText(type="output_text", text="one", annotations=[]), + OutputText(type="output_text", text="two", annotations=[]), + ], + ) + key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])]) + assert key.texts == ("one", "two") + + def test_output_item_done_round_is_never_deduped(self): + done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} + assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index befd4c5ffbd..4cff5c76b9e 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -7,12 +7,13 @@ transformations for the Responses API. Source: litellm/llms/xai/responses/transformation.py """ -from unittest.mock import MagicMock - +from unittest.mock import MagicMock, Mock +import httpx import pytest import litellm +from litellm.llms.xai.cost_calculator import cost_per_token from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import ( @@ -400,3 +401,94 @@ class TestXAIResponsesWebSearchBilling: bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS + + +class TestXAIResponsesReportedCost: + """xAI reports what it charged; the transformation moves it to where litellm bills from. + + ``ResponseAPILoggingUtils`` copies ``usage.cost`` onto the chat Usage that cost + tracking prices, so restating ``cost_in_usd_ticks`` there is what makes /v1/responses + bill the reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756. + """ + + @staticmethod + def _response_body(usage: dict) -> dict: + return { + "id": "resp_xai", + "object": "response", + "created_at": 0, + "model": "grok-4-latest", + "status": "completed", + "output": [], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": usage, + } + + def _transformed_usage(self, usage: dict) -> ResponseAPIUsage | None: + raw_response = httpx.Response(status_code=200, json=self._response_body(usage)) + + response = XAIResponsesAPIConfig().transform_response_api_response( + model="grok-4-latest", + raw_response=raw_response, + logging_obj=Mock(), + ) + return response.usage + + def test_reported_cost_reaches_the_cost_calculator(self): + usage = self._transformed_usage( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ) + + assert usage.cost == 0.0037756 + + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + + def test_streamed_reported_cost_reaches_the_cost_calculator(self): + event = XAIResponsesAPIConfig().transform_streaming_response( + model="grok-4-latest", + parsed_chunk={ + "type": "response.completed", + "sequence_number": 7, + "response": self._response_body( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ), + }, + logging_obj=Mock(), + ) + + assert isinstance(event, ResponseCompletedEvent) + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) + assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + + def test_usage_without_a_reported_cost_is_left_alone(self): + usage = self._transformed_usage( + {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300} + ) + + assert usage.cost is None + + def test_negative_reported_cost_is_not_carried(self): + """A caller who can set api_base must not be able to report negative spend.""" + usage = self._transformed_usage( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": -37756000, + } + ) + + assert usage.cost is None diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index e5e853ec82f..c67dca11a56 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -1,9 +1,14 @@ +from unittest.mock import Mock - +import httpx import pytest import litellm -from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.chat.transformation import ( + XAIChatCompletionStreamingHandler, + XAIChatConfig, +) +from litellm.llms.xai.cost_calculator import cost_per_token from litellm.types.utils import ( CompletionTokensDetailsWrapper, ModelResponse, @@ -195,3 +200,113 @@ class TestXAIChatWebSearchBilling: ) assert with_search - without_search == pytest.approx(3 * 5.0 / 1000.0) + + +class TestXAIReportedCost: + """xAI reports what it charged; the transformation moves it to where litellm bills from. + + ``cost`` is the field litellm already carries a provider stated cost in, so restating + ``cost_in_usd_ticks`` there is what lets ``llms/xai/cost_calculator.py`` bill the + reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756. + """ + + @staticmethod + def _transformed_usage(usage: dict) -> Usage: + raw_response = httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-xai", + "object": "chat.completion", + "created": 0, + "model": "grok-4-latest", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": usage, + }, + ) + + response = XAIChatConfig().transform_response( + model="grok-4-latest", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + return response.usage + + def test_reported_cost_reaches_the_cost_calculator(self): + usage = self._transformed_usage( + { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ) + + assert usage.cost == 0.0037756 + assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0037756) + + def test_usage_without_a_reported_cost_is_left_alone(self): + usage = self._transformed_usage( + {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300} + ) + + assert getattr(usage, "cost", None) is None + + def test_negative_reported_cost_is_not_carried(self): + """A caller who can set api_base must not be able to report negative spend.""" + usage = self._transformed_usage( + { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": -37756000, + } + ) + + assert getattr(usage, "cost", None) is None + + def test_streamed_reported_cost_survives_chunk_aggregation(self): + """Streamed spend only matches if the conversion happens on the chunk. + + Chunk aggregation rebuilds usage from the fields it models plus ``cost``, so a + chunk still carrying only ``cost_in_usd_ticks`` loses the reported amount. + """ + handler = XAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + + parsed = handler.chunk_parser( + { + "id": "chatcmpl-xai", + "object": "chat.completion.chunk", + "created": 0, + "model": "grok-4-latest", + "choices": [], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + }, + } + ) + + assert parsed.usage.cost == 0.0037756 + + assembled = litellm.stream_chunk_builder(chunks=[parsed]) + assert assembled.usage.cost == 0.0037756 + assert cost_per_token(model="grok-4-latest", usage=assembled.usage) == ( + 0.0, + 0.0037756, + ) diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 92e76fd18ab..6503e956a51 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -7,7 +7,10 @@ import os import litellm from litellm.types.utils import ( + Choices, CompletionTokensDetailsWrapper, + Message, + ModelResponse, PromptTokensDetailsWrapper, Usage, ) @@ -361,6 +364,145 @@ class TestXAICostCalculator: response_object=object(), usage=usage ) + def test_reported_cost_is_preferred_over_token_math(self): + """The amount xAI reported, carried on usage.cost by the transformation, is billed. + + It lands entirely on completion cost because xAI does not split its total by + direction, the same shape the perplexity calculator returns. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=0.0037756, + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost == 0.0 + assert math.isclose(completion_cost, 0.0037756, rel_tol=1e-10) + + def test_reported_cost_suppresses_web_search_surcharge(self): + """The reported total already covers server-side tool calls. + + Without the suppression these 3 searches would be billed a second time on + top of the total xAI already charged. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=100, + web_search_requests=3, + ), + cost=0.0037756, + ) + + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_web_search_surcharge_suppressed_through_the_dispatcher(self): + """The suppression has to hold on the path cost tracking actually uses. + + Legacy behaviour stays intact when xAI reports no cost. + """ + from litellm.llms import get_cost_for_web_search_request + + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + assert get_cost_for_web_search_request("xai", usage, {}) > 0.0 + + reported = Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150, cost=0.0037756 + ) + setattr(reported, "server_side_tool_usage_details", {"web_search_calls": 3}) + assert get_cost_for_web_search_request("xai", reported, {}) == 0.0 + + def test_no_reported_cost_falls_back_to_token_math(self): + """Absent the provider figure, nothing changes for existing callers.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + + def test_malformed_reported_cost_falls_back_to_token_math(self): + """A junk value must not fail the request, fall back to calculating.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + setattr(usage, "cost", "not-a-number") + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + + def test_boolean_reported_cost_falls_back_to_token_math(self): + """True is an int in python and would otherwise be billed as $1.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + setattr(usage, "cost", True) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + assert completion_cost != 1.0 + + def test_negative_reported_cost_is_rejected(self): + """A negative amount must never reach spend tracking. + + A caller who can set api_base controls the response body, so trusting a + negative figure would let them subtract from their own recorded spend and + slip past a budget. Fall back to token pricing instead, and keep charging + the web search surcharge, since no trustworthy total was reported. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=-0.0037756, + ) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0 + + def test_non_finite_reported_cost_is_rejected(self): + """NaN compares false against every budget threshold. + + Usage stores a provider supplied cost without validating it, so a caller who + controls the response body could report NaN and leave spend >= max_budget + false for the life of the key rather than mispricing one request. The + infinities are refused alongside it. Fall back to token pricing and keep + charging the web search surcharge, since no trustworthy total was reported. + """ + for reported_cost in (float("nan"), float("inf"), float("-inf")): + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=reported_cost, + ) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert math.isfinite(prompt_cost), reported_cost + assert math.isfinite(completion_cost), reported_cost + assert prompt_cost > 0.0, reported_cost + assert completion_cost > 0.0, reported_cost + assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0, reported_cost + + def test_zero_reported_cost_is_honoured(self): + """A reported zero is a real answer, not a missing value.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300, cost=0.0) + + assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0) + def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) @@ -437,6 +579,48 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + def test_custom_pricing_beats_the_reported_cost(self): + response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")], + usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956), + ) + + billed = litellm.completion_cost( + completion_response=response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.001}, + custom_pricing=True, + ) + + assert math.isclose(billed, 0.551, rel_tol=1e-10) + + def test_deployment_custom_pricing_beats_the_reported_cost(self, monkeypatch): + deployment_id = "xai-deployment-priced-by-the-operator" + monkeypatch.setitem( + litellm.model_cost, + deployment_id, + {"input_cost_per_token": 0.001, "output_cost_per_token": 0.001, "litellm_provider": "xai", "mode": "chat"}, + ) + response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")], + usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956), + ) + + billed = litellm.completion_cost( + completion_response=response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + custom_pricing=True, + router_model_id=deployment_id, + ) + + assert math.isclose(billed, 0.551, rel_tol=1e-10) + class TestXAIWebSearchCostHelpers: """Focused coverage for web_search / tool-usage helpers in cost_calculator.py.""" 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/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 476d443d8d8..cb6772977ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None): @pytest.mark.asyncio async def test_openai_moderation_streaming_default_uses_sampled_cadence(): """Default config samples every 5th streamed chunk and runs a final aggregate - pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream - calls, plus 1 final = 3 total. + pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream + round is skipped because chunk 10 already scanned the full text, for 2 total calls """ import litellm @@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence(): ): pass - assert patched_make_request.await_count == 3, ( - f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), " + assert patched_make_request.await_count == 2, ( + f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {patched_make_request.await_count}" ) @@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder @pytest.mark.asyncio async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled(): """With streaming_end_of_stream_only=False and streaming_sampling_rate=2, - moderation runs every 2nd chunk during the stream, plus once more at end. + moderation runs every 2nd chunk during the stream. The terminal chunk scan covers + the final aggregate, for 3 total calls """ import litellm @@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab ): pass - # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls), - # plus the final aggregate pass after the stream ends (1 call) = 4 total. - assert patched_make_request.await_count == 4, ( - f"Expected 4 moderation calls (3 sampled + 1 final aggregate), " + assert patched_make_request.await_count == 3, ( + f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {patched_make_request.await_count}" ) 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_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 523ec1a37b4..83cc9ae8bb9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_default_uses_sampled_cadence(self): - """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3.""" + """Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10 + cover the full text, so the end-of-stream round is skipped and there are 2 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 3, ( - f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), " + assert mock_post.await_count == 2, ( + f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {mock_post.await_count}" ) for call in mock_post.await_args_list: @@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_sampling_rate_override(self): - """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls.""" + """sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so + the end-of-stream round is skipped and there are 3 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 4, ( - f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), " + assert mock_post.await_count == 3, ( + f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {mock_post.await_count}" ) 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/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8cad1c634a9..a28a2a71613 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket: assert recorded[0]["guardrail_name"] == "audit-recorder" assert recorded[0]["guardrail_status"] == "success" assert request_data["metadata"]["user_api_key_user_id"] == "user-1" + + +class _ScanCountingGuardrail(CustomGuardrail): + """Pass-through guardrail that records every response-side scan payload.""" + + def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False): + super().__init__(guardrail_name="scan-counter") + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = end_of_stream_only + self.streaming_buffer_until_moderated = buffer_until_moderated + self.guardrail_config = {} + self.scans: tuple[dict[str, object], ...] = () + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.scans = ( + *self.scans, + { + "texts": list(inputs.get("texts") or []), + "tool_calls": list(inputs.get("tool_calls") or []), + "model": inputs.get("model"), + }, + ) + return inputs + + +def _responses_delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + +def _responses_tail(sequence_number, text): + return [ + { + "type": "response.output_text.done", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": text, + }, + { + "type": "response.completed", + "sequence_number": sequence_number + 1, + "response": { + "model": "gpt-5.6", + "output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}], + }, + }, + ] + + +class TestStreamingScanDedup: + """A sampled round whose scan payload matches the previous round (or carries + no text yet) is skipped, so a stream is never re-scanned for output the + guardrail already cleared. Regression for LIT-6692.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self, monkeypatch): + monkeypatch.setattr( + unified_module, + "endpoint_guardrail_translation_mappings", + load_guardrail_translation_mappings(), + ) + + @pytest.mark.asyncio + async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 3 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a"), + _stream_chunk("b"), + _stream_chunk("c"), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk("d", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self): + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + guardrail = _ScanCountingGuardrail(sampling_rate=3) + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + finish = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls") + ] + ) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a", index=0), + _stream_chunk("x", index=1), + _stream_chunk("b", finish_reason="stop", index=0), + _stream_chunk("y", index=1), + _stream_chunk("z", finish_reason="stop", index=1), + ] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(guardrail.scans) == 2 + assert any("yz" in text for text in guardrail.scans[-1]["texts"]) + + @pytest.mark.asyncio + async def test_responses_completed_event_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(8)] + full_text = "".join(f"t{i}" for i in range(8)) + chunks = deltas + _responses_tail(8, full_text) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 10 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]] + assert guardrail.scans[-1]["model"] == "gpt-5.6" + + @pytest.mark.asyncio + async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + chunks = deltas + _responses_tail(5, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]] + + @pytest.mark.asyncio + async def test_responses_completed_carrying_a_function_call_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + completed = { + "type": "response.completed", + "sequence_number": 5, + "response": { + "model": "gpt-5.6", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + }, + } + chunks = deltas + [completed] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 6 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_responses_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)] + chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 12 + assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},) + + @pytest.mark.asyncio + async def test_responses_tool_call_done_event_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + tool_call_done = { + "type": "response.output_item.done", + "sequence_number": 1, + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + } + chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 4 + assert len(guardrail.scans) == 2 + assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"] + assert guardrail.scans[1]["texts"] == ["hi"] + + @pytest.mark.asyncio + async def test_anthropic_skips_empty_round_and_terminal_duplicate(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_still_scans_exactly_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self): + guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index d85c6da659b..b99dcb062b4 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -670,6 +670,37 @@ def test_get_provider_specific_params(): ) # Literal type should be select +@pytest.mark.asyncio +async def test_provider_specific_params_includes_hide_secrets(): + """hide-secrets lives in the enterprise package so it is not in + guardrail_class_registry; the endpoint must still advertise it or the + Add Guardrail UI dropdown never offers it (LIT-3548).""" + from litellm.proxy.guardrails.guardrail_endpoints import ( + get_provider_specific_params, + ) + + provider_params = await get_provider_specific_params() + + assert "hide-secrets" in provider_params + # populateGuardrailProviders() in the dashboard only lists providers whose + # entry carries a ui_friendly_name. + assert provider_params["hide-secrets"]["ui_friendly_name"] == "Hide Secrets" + assert provider_params["hide-secrets"]["detect_secrets_config"]["required"] is False + + +@pytest.mark.asyncio +async def test_add_guardrail_settings_restricts_hide_secrets_to_pre_call(): + """hide-secrets only implements async_pre_call_hook, so offering the other + modes in the UI would create configs that boot clean and never run.""" + from litellm.proxy.guardrails.guardrail_endpoints import ( + get_guardrail_ui_settings, + ) + + settings = await get_guardrail_ui_settings() + + assert settings.supported_modes_by_provider["hide-secrets"] == ["pre_call"] + + def test_optional_params_not_returned_when_not_overridden(): """Test that optional_params is not returned when the config model doesn't override it""" from typing import Optional diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 36d6414ba9f..635e6332958 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -516,6 +516,39 @@ async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated): assert result.team_ids == ["team-a", "team-b"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "roles", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +async def test_multi_valued_role_attribute_resolves_to_highest_privilege(saml_env_idp_initiated, roles): + """An assertion carrying several roles must not depend on the order the IdP emitted them in.""" + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + attributes={ + "email": ["dave@example.com"], + "role": roles, + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_assertion_without_role_attribute_has_no_user_role(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, attributes={"email": ["erin@example.com"]}) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role is None + + @pytest.mark.asyncio async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): cache = DualCache() diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index dd8c752a868..5dfff53f7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -6598,13 +6598,94 @@ def test_get_litellm_user_role_with_invalid_role(): assert result is None -def test_get_litellm_user_role_with_list_multiple_roles(): - """Test that get_litellm_user_role takes the first element from a multi-element list.""" +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin", "internal_user"], + ["internal_user", "proxy_admin"], + ], +) +def test_get_litellm_user_role_picks_highest_privilege_regardless_of_order(role_claim): + """A multi-valued role claim resolves to the most privileged role, not the first one listed.""" from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.types import get_litellm_user_role - result = get_litellm_user_role(["proxy_admin", "internal_user"]) - assert result == LitellmUserRoles.PROXY_ADMIN + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin_viewer", "internal_user"], + ["internal_user", "proxy_admin_viewer"], + ], +) +def test_get_litellm_user_role_keeps_org_spend_visibility_for_mixed_roles(role_claim): + """ + Regression for LIT-6077: a user holding both proxy_admin_viewer and internal_user kept + losing org-level spend visibility whenever the IdP happened to list internal_user first. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +def test_get_litellm_user_role_ignores_unrecognised_entries(): + """Roles LiteLLM does not know about are skipped rather than swallowing the whole claim.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["some_idp_group", "internal_user"]) == LitellmUserRoles.INTERNAL_USER + assert get_litellm_user_role(["some_idp_group", "another_group"]) is None + + +def test_get_litellm_user_role_list_lookup_is_case_insensitive(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["INTERNAL_USER", "Proxy_Admin"]) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "team"], + ["team", "org_admin"], + ], +) +def test_get_litellm_user_role_is_deterministic_for_unranked_roles(role_claim): + """Roles outside the privilege hierarchy still resolve the same way in either claim order.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.ORG_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "internal_user"], + ["internal_user", "org_admin"], + ], +) +def test_get_litellm_user_role_prefers_a_ranked_role_over_an_unranked_one(role_claim): + """ + org_admin, team and customer sit outside the privilege ladder, so a claim mixing one of + them with a ranked role settles on the ranked role in either order. Same rule the Entra + app_roles and role_mappings paths already follow. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.INTERNAL_USER + + +def test_get_litellm_user_role_returns_none_for_non_string_claims(): + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(None) is None + assert get_litellm_user_role({"role": "proxy_admin"}) is None # ============================================================================ @@ -6654,6 +6735,46 @@ def test_process_sso_jwt_access_token_extracts_role_from_access_token(): assert result.user_role == LitellmUserRoles.PROXY_ADMIN +@pytest.mark.parametrize( + "role_claim", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +def test_process_sso_jwt_access_token_resolves_highest_privilege_role(role_claim): + """ + The generic SSO access-token path must land on the same role for a user whose role + claim holds several roles, whichever order the IdP emitted them in. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + access_token_str = pyjwt.encode( + {"sub": "user-123", "email": "mixed@test.com", "litellm_role": role_claim}, + "secret", + algorithm="HS256", + ) + result = CustomOpenID( + id="user-123", + email="mixed@test.com", + display_name="Mixed Role User", + team_ids=[], + user_role=None, + ) + + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + def test_process_sso_jwt_access_token_does_not_override_existing_role(): """ Test that process_sso_jwt_access_token does NOT override a role that was diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 551e27a18f5..41439f28638 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1104,6 +1104,16 @@ def test_get_autorouter_presets_local_mode_serves_bundled_catalog( assert response.status_code == 200 payload = response.json() assert "anthropic_family" in payload + assert payload["1m_context"]["complexity_router_config"]["classifier_type"] == "heuristic_v2" + assert payload["1m_context"]["complexity_router_config"]["tiers"] == { + "SIMPLE": ["gpt-5.6-luna"], + "MEDIUM": ["gpt-5.6-terra"], + "COMPLEX": ["claude-opus-5"], + "REASONING": ["claude-opus-5"], + } + assert payload["1m_context"]["complexity_router_config"]["tier_model_configs"] == { + "REASONING": [{"model_name": "claude-opus-5", "litellm_params": {"reasoning_effort": "high"}}] + } for preset in payload.values(): assert isinstance(preset["label"], str) assert isinstance(preset["description"], str) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index ee0e2014951..8366e5546a9 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem @@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): } assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + + +def _request_for(path: str) -> MagicMock: + request = MagicMock(spec=Request) + request.scope = {"path": path} + request.url = MagicMock() + request.url.path = path + request.url.__str__.return_value = f"http://localhost{path}" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + return request + + +def _spend_log_session_id(data: dict[str, object]) -> str: + """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log + + metadata = data["metadata"] + assert isinstance(metadata, dict) + litellm_params = get_litellm_params( + litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None, + litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None, + metadata=metadata, + ) + trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"), + litellm_params=litellm_params, + ) + return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_correlation_in_logs", [False, True]) +async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree( + monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool +): + """Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs + must be the same generated id, so cross-referencing the two by session_id works. The id is marked + as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs) + data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} + + updated = await add_litellm_data_to_request( + data=data, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + callback_session_id = updated["metadata"]["session_id"] + assert isinstance(callback_session_id, str) and len(callback_session_id) == 36 + assert _spend_log_session_id(updated) == callback_session_id + assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True + assert get_fireworks_session_id( + {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]} + ) is None + + +@pytest.mark.asyncio +async def test_missing_session_id_unset_keeps_legacy_divergence(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + + assert "session_id" not in updated["metadata"] + assert "litellm_session_id" not in updated + assert _spend_log_session_id(updated) == "per-call-random-trace-id" + + +@pytest.mark.asyncio +async def test_missing_session_id_generate_reuses_traceparent_trace_id(): + """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" + request = _request_for("/v1/chat/completions") + request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["generate", "reject"]) +async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str): + request = _request_for("/v1/chat/completions") + request.headers = {"x-litellm-session-id": "client-session-1"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": policy}, + ) + + assert updated["litellm_session_id"] == "client-session-1" + assert updated["metadata"]["session_id"] == "client-session-1" + assert _spend_log_session_id(updated) == "client-session-1" + assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"] + assert ( + get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]}) + == "client-session-1" + ) + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_accepts_body_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert updated["metadata"]["session_id"] == "body-session-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_returns_400_without_session_id(): + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "session_id" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"]) +async def test_missing_session_id_policy_skips_non_inference_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert "session_id" not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_missing_session_id_unknown_value_is_ignored(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "typo"}, + ) + + assert "session_id" not in updated["metadata"] diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 941d78085e4..3ecccb673f3 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -16,7 +16,7 @@ import litellm from litellm import Router from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -4274,6 +4274,26 @@ class TestSessionAffinity: assert first.model == "o1-preview" assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio + async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config): + """A session id the proxy generated for a request that had none is per request, so + it must not create a pin even with session_affinity enabled.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}} + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the @@ -10777,6 +10797,9 @@ class TestModalityRouting: ("custom_tiers_walk", "premium-model", "modality_escalation"), ("pin_kept_bypasses", "text-cheap", "session_affinity_pin"), ("pin_replacement_gated", "vision-big", "modality_escalation"), + ("pin_override_escalates", "vision-mid", "modality_pin_override"), + ("pin_override_same_tier", "vision-cheap", "modality_pin_override"), + ("pin_override_inert_without_modality_routing", "text-cheap", "session_affinity_pin"), ("adaptive_pick_rewritten", "vision-mid", "modality_escalation"), ], ) @@ -10827,7 +10850,7 @@ class TestModalityRouting: messages = [ {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]} ] - elif path in ("pin_kept_bypasses", "pin_replacement_gated"): + elif path.startswith(("pin_kept", "pin_replacement", "pin_override")): cache = AsyncMock() cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache @@ -10839,6 +10862,13 @@ class TestModalityRouting: messages = [ {"role": "user", "content": [{"type": "text", "text": "LITELLM ESCALATE describe this"}, IMG_PART]} ] + elif path == "pin_override_same_tier": + config["modality_pin_override"] = True + config["tiers"]["SIMPLE"] = ["text-cheap", "vision-cheap"] + vision["vision-cheap"] = True + elif path == "pin_override_inert_without_modality_routing": + config["modality_routing"] = False + config["modality_pin_override"] = path.startswith("pin_override") elif path == "adaptive_pick_rewritten": config["adaptive"] = True mock_router_instance.model_list = [] @@ -11005,4 +11035,60 @@ class TestModalityRouting: from litellm.router_strategy.complexity_router.complexity_router import _decision_is_pinnable assert _decision_is_pinnable({"cause": "modality_escalation"}) is False + assert _decision_is_pinnable({"cause": "modality_pin_override"}) is False assert _decision_is_pinnable({"cause": "heuristic_scorer"}) is True + + @pytest.mark.asyncio + async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance): + """The override is for one request: the session keeps the model it was pinned to.""" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) + mock_router_instance.cache = cache + router = self._router( + mock_router_instance, + { + "tiers": dict(self.BASE_TIERS), + "modality_routing": True, + "modality_pin_override": True, + "session_affinity": True, + }, + dict(self.BASE_VISION), + ) + request_kwargs = {"metadata": {"session_id": "s1"}} + + image_turn = await router.async_pre_routing_hook( + model="m", request_kwargs=request_kwargs, messages=self.IMAGE_MESSAGE + ) + assert image_turn.model == "vision-mid" + assert image_turn.routing_decision["cause"] == "modality_pin_override" + assert "modality_escalated_from:SIMPLE" in image_turn.routing_decision["signals"] + + assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"} + + text_turn = await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=[{"role": "user", "content": "hi"}] + ) + assert text_turn.model == "text-cheap" + assert text_turn.routing_decision["cause"] == "session_affinity_pin" + + @pytest.mark.asyncio + async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance): + """The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin.""" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) + mock_router_instance.cache = cache + router = self._router( + mock_router_instance, + { + "tiers": {"SIMPLE": "text-cheap", "COMPLEX": "text-big"}, + "modality_routing": True, + "modality_pin_override": True, + "session_affinity": True, + }, + {"text-cheap": False, "text-big": False}, + ) + with pytest.raises(litellm.BadRequestError, match="no model"): + await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE + ) + assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"} diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index a3772a276fa..cf48888600e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -7,7 +7,7 @@ import json import litellm from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert filtered[0]["model_info"]["id"] == "deployment-2" +@pytest.mark.asyncio +async def test_proxy_generated_session_id_does_not_pin_a_deployment(): + """A session id the proxy generated for a request that had none is per request, so a + pin stored under it must be ignored and none must be written.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + healthy_deployments = [ + {"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}}, + {"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}}, + ] + await cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"), + {"model_id": "deployment-2"}, + ) + request_kwargs = { + "metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True} + } + + filtered = await callback.async_filter_deployments( + model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs + ) + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"}, + "model_info": {"id": "deployment-1"}, + }, + call_type=None, + ) + + assert len(filtered) == 2 + assert await cache.async_get_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1") + ) == {"model_id": "deployment-2"} + + MOCK_RESPONSES_API_RESPONSE = { "id": "resp_mock-resp-456", "object": "response", diff --git a/tests/test_litellm/test_check_py310_typing_imports.py b/tests/test_litellm/test_check_py310_typing_imports.py new file mode 100644 index 00000000000..de370326091 --- /dev/null +++ b/tests/test_litellm/test_check_py310_typing_imports.py @@ -0,0 +1,86 @@ +import sys +from pathlib import Path +from typing import Final + +_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests" +sys.path.insert(0, str(_CODE_COVERAGE_DIR)) # test-quality-ok: required to import checker from its source directory +import check_py310_typing_imports as checker # noqa: E402 # load checker from its source directory + + +def _scan(tmp_path: Path, source: str) -> tuple[object, ...]: + file_path = tmp_path / "fixture.py" + file_path.write_text(source, encoding="utf-8") + return checker.scan_file(file_path) + + +def test_typing_import_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n") + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_typing_extensions_import_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == () + + +def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "import typing\nx: typing.Self\n") + assert tuple(violation.name for violation in violations) == ("Self",) + + +def test_version_guarded_typing_import_passes(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing import NotRequired\n" + "else:\n" + " from typing_extensions import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + violations = _scan(tmp_path, source) + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info < (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_nearest_if_controls_version_guard(tmp_path: Path) -> None: + source = ( + "if sys.version_info >= (3, 11):\n" + " from typing import Self\n" + " x = 1\n" + "if True:\n" + " from typing import Self\n" + ) + violations = _scan(tmp_path, source) + assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),) + + +def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None: + file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py" + file_path.parent.mkdir(parents=True) + file_path.write_text("from typing import NotRequired\n", encoding="utf-8") + + violations = checker.scan_directory(tmp_path) + assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),) + + +def test_python_310_typing_name_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing import Optional\n") == () diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 7c2b9d0be05..3c8bf142835 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3131,9 +3131,9 @@ def test_stream_chunk_builder_prices_proxy_alias_via_model_map(): assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) -def _stream_builder_logging_obj() -> LiteLLMLogging: +def _stream_builder_logging_obj(model: str = "gpt-4o", custom_llm_provider: str = "openai") -> LiteLLMLogging: logging_obj: Final = LiteLLMLogging( - model="gpt-4o", + model=model, messages=[{"role": "user", "content": "hi"}], stream=True, call_type="completion", @@ -3142,10 +3142,11 @@ def _stream_builder_logging_obj() -> LiteLLMLogging: function_id="test-function-id", ) logging_obj.update_environment_variables( - model="gpt-4o", + model=model, user=None, optional_params={}, - litellm_params={"custom_llm_provider": "openai"}, + litellm_params={"custom_llm_provider": custom_llm_provider}, + custom_llm_provider=custom_llm_provider, ) return logging_obj @@ -3237,3 +3238,24 @@ def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): assert response.usage.completion_tokens == 60 assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) assert response._hidden_params["response_cost"] == pytest.approx(0.000704) + + +def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + usage_chunk: Final = _stream_builder_text_chunk("grok-4", "") + usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) + chunks: Final = [ + _stream_builder_text_chunk("grok-4", "Hello "), + _stream_builder_text_chunk("grok-4", "world.", finish_reason="stop"), + usage_chunk, + ] + logging_obj: Final = _stream_builder_logging_obj(model="grok-4", custom_llm_provider="xai") + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(0.42) + assert response._hidden_params.get("response_cost") is None + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) 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..d5bf3883be4 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22334 + "limit": 22330 }, "LIT002": { "limit": 26763 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16480 + "limit": 16477 }, "LIT011": { - "limit": 5520 + "limit": 5519 }, "LIT012": { "limit": 4489 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7de7373b20b..7d7066addcf 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1399,9 +1399,6 @@ }, "prefer-const": { "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/components/TeamsPage/teamTableColumns.tsx": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index c1f2ddcf51c..f686ff5644a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -201,6 +201,7 @@ export const guardrailLogoMap = { XecGuard: xecguardLogo.src, "LiteLLM Content Filter": litellmLogo.src, "LiteLLM LLM as a Judge": litellmLogo.src, + "Hide Secrets": litellmLogo.src, Akto: aktoLogo.src, "DeepKeep AI Firewall": deepkeepLogo.src, "Qostodian Nexus": qohashLogo.src, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.test.tsx new file mode 100644 index 00000000000..ab63b7637b1 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.test.tsx @@ -0,0 +1,99 @@ +import React from "react"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import { useForm } from "react-hook-form"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import GuardrailProviderFields from "./guardrail_provider_fields"; +import { populateGuardrailProviderMap } from "./guardrail_info_helpers"; +import type { GuardrailFormValues } from "./GuardrailFormField"; + +vi.mock("@/lib/toast", () => ({ toast: { error: vi.fn() } })); + +const HIDE_SECRETS_PARAMS = { + "hide-secrets": { + ui_friendly_name: "Hide Secrets", + detect_secrets_config: { + param: "detect_secrets_config", + description: "Optional detect-secrets configuration", + required: false, + type: "object", + }, + }, +}; + +const Harness: React.FC<{ onValid: (values: GuardrailFormValues) => void }> = ({ onValid }) => { + const form = useForm(); + return ( +
+ + + + ); +}; + +const renderHarness = () => { + populateGuardrailProviderMap(HIDE_SECRETS_PARAMS); + const onValid = vi.fn(); + renderWithProviders(); + const textarea = screen.getByLabelText(/detect_secrets_config/) as HTMLTextAreaElement; + return { onValid, textarea }; +}; + +describe("GuardrailProviderFields object field", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("commits a valid JSON object as a parsed dict", async () => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: '{"plugins_used": [{"name": "AWSKeyDetector"}]}' } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await waitFor(() => expect(onValid).toHaveBeenCalledTimes(1)); + expect(onValid.mock.calls[0][0].detect_secrets_config).toEqual({ + plugins_used: [{ name: "AWSKeyDetector" }], + }); + }); + + it("blocks submission while the field holds malformed JSON", async () => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: "{not json" } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await screen.findByText("detect_secrets_config must be a valid JSON object"); + expect(onValid).not.toHaveBeenCalled(); + expect(textarea.value).toBe("{not json"); + }); + + it.each(['["array"]', '"scalar"', "null", "42"])("blocks non-object JSON %s", async (raw) => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: raw } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await screen.findByText("detect_secrets_config must be a valid JSON object"); + expect(onValid).not.toHaveBeenCalled(); + }); + + it("treats a cleared field as unset and submits", async () => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: '{"a": 1}' } }); + fireEvent.blur(textarea); + fireEvent.change(textarea, { target: { value: "" } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await waitFor(() => expect(onValid).toHaveBeenCalledTimes(1)); + expect(onValid.mock.calls[0][0].detect_secrets_config).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx index 76c91d4c176..7d48da89ebf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx @@ -13,6 +13,8 @@ import { FieldGroup } from "@/components/ui/field"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; +import { Textarea } from "@/components/ui/textarea"; +import { toast } from "@/lib/toast"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { asStringArray, @@ -22,6 +24,7 @@ import { readRecord, requiredRule, type GuardrailFieldControlProps, + type GuardrailFieldRules, type GuardrailFormControl, } from "./GuardrailFormField"; @@ -60,6 +63,44 @@ const BOOLEAN_ITEMS = [ const isSecretKey = (fieldKey: string): boolean => fieldKey.includes("password") || fieldKey.includes("secret") || fieldKey.includes("key"); +const isPlainObject = (value: unknown): value is Record => + typeof value === "object" && value !== null && !Array.isArray(value); + +// Object fields hold the raw text while the user types, so submission must be +// blocked until the value parses to a plain JSON object (or is cleared). +const jsonObjectRule = (fieldKey: string): GuardrailFieldRules => ({ + validate: (value: unknown) => + value === undefined || isPlainObject(value) ? true : `${fieldKey} must be a valid JSON object`, +}); + +// Commits a parsed object (or undefined for a cleared field) to the form on +// blur; anything else stays as raw text so jsonObjectRule blocks submission. +const commitObjectField = (raw: string, onChange: (value: unknown) => void): void => { + const next = raw.trim(); + if (next === "") { + onChange(undefined); + return; + } + let parsed: unknown; + try { + parsed = JSON.parse(next); + } catch { + parsed = next; + } + if (isPlainObject(parsed)) { + onChange(parsed); + } else { + toast.error("Enter a valid JSON object for this configuration"); + } +}; + +const fieldRules = (field: ProviderParam, fieldKey: string): GuardrailFieldRules | undefined => { + if (field.type === "object") { + return jsonObjectRule(fieldKey); + } + return field.required ? requiredRule(`${fieldKey} is required`) : undefined; +}; + interface ProviderFieldInputProps { descriptor: ProviderParam; fieldKey: string; @@ -141,6 +182,25 @@ const ProviderFieldInput: React.FC = ({ descriptor, fie ); } + if (descriptor.type === "object") { + const objectValue = typeof value === "object" && value !== null ? JSON.stringify(value, null, 2) : asText(value); + return ( +