diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml index 7e40a860ee9..838ba35f831 100644 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ b/.github/workflows/auto_update_price_and_context_window.yml @@ -24,12 +24,9 @@ jobs: - name: Update JSON Data run: | uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py" - - name: Regenerate JSON Schema - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py - name: Create Pull Request run: | - git add model_prices_and_context_window.json model_prices_and_context_window.schema.json + git add model_prices_and_context_window.json git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')" gh pr create --title "Update model_prices_and_context_window.json file" \ --body "Automated update for model_prices_and_context_window.json" \ diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml index 18d5e3de1eb..5325510e45e 100644 --- a/.github/workflows/sync-together-ai-models.yml +++ b/.github/workflows/sync-together-ai-models.yml @@ -41,10 +41,6 @@ jobs: uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md" env: TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }} - - name: Regenerate the JSON schema - if: steps.existing.outputs.open_pr == '' - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py - name: Create a pull request when the registry changed if: steps.existing.outputs.open_pr == '' run: | @@ -57,8 +53,7 @@ jobs: git config user.email "41898282+github-actions[bot]@users.noreply.github.com" git checkout -b "$branch" git add model_prices_and_context_window.json \ - litellm/model_prices_and_context_window_backup.json \ - model_prices_and_context_window.schema.json + litellm/model_prices_and_context_window_backup.json git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')" gh auth setup-git git push origin "$branch" diff --git a/README.md b/README.md index 3bd0fd54a65..77f7e3baa4d 100644 --- a/README.md +++ b/README.md @@ -384,6 +384,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | | | [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | | | [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | | +| [Nace (`nace`)](https://console.nace.ai/docs/drex) | | | | | | | | | | | | [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ | | | | | | | | | [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/backend/Dockerfile b/backend/Dockerfile index 5213bb00c50..84c90e16f67 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -4,6 +4,33 @@ ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2 FROM $UV_IMAGE AS uvbin +# ---------- Rust extension ---------- +# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes +FROM $LITELLM_BUILD_IMAGE AS rust-native + +WORKDIR /app +USER root + +COPY --from=uvbin /uv /uvx /usr/local/bin/ + +RUN for i in 1 2 3; do \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done + +ENV UV_PYTHON_DOWNLOADS=0 + +COPY pyproject.toml README.md rust-toolchain.toml ./ +COPY .cargo ./.cargo +COPY litellm-rust ./litellm-rust + +RUN --mount=type=cache,target=/root/.cache/uv \ + mkdir -p litellm/rust_bridge /wheels && \ + touch litellm/__init__.py litellm/rust_bridge/__init__.py && \ + uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ + python -c "import maturin; maturin.build_wheel('/wheels')" + # ---------- Builder ---------- FROM $LITELLM_BUILD_IMAGE AS builder @@ -54,7 +81,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ COPY . . RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-default-groups --no-editable \ + uv sync --frozen --no-install-project --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --group admin-mcp \ @@ -63,6 +90,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra saml \ --python python3.13 +RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \ + uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl + RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 60d695653b5..2a7522cf34e 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -165,6 +165,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset( BACKEND_MOUNT_PATHS: frozenset[str] = frozenset( { "/admin", + "/metrics", "/swagger", # API documentation static assets belong to the backend "/mcp", # lazily-mounted MCP sub-app serves on the backend component } diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index e7b873e3f5c..932da84e5d7 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -1,11 +1,11 @@ """Guard the cost map on pull requests. -Every pull request whose diff against its merge base touches one of the three cost map files gets the file -checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the -map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies -and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot -(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three -files and may only add or update models. +Every pull request whose diff against its merge base touches one of the two cost map files gets the file +checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched +skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the +branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always +get the file checks and additionally may only touch those two files or the legacy schema +path and may only add or update models. """ from __future__ import annotations @@ -19,13 +19,12 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final -from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors - COST_MAP_PATH: Final = "model_prices_and_context_window.json" BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" -SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" -GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH) +LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" +GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH) BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" +SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"}) EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({}) CostMap = dict[str, object] @@ -35,7 +34,6 @@ CostMap = dict[str, object] class Snapshot: cost_map: str backup: str - schema: str def _parse_object(text: str, path: str) -> CostMap | str: @@ -46,13 +44,6 @@ def _parse_object(text: str, path: str) -> CostMap | str: return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root" -def _rendered_schema(cost_map: CostMap) -> str: - try: - return render(build_schema(cost_map)) - except SystemExit as error: - return str(error) - - def _evaluation_output_price_failure( key: str, entry: object, evaluation_output_price_sources: Mapping[str, str] ) -> str | None: @@ -78,29 +69,13 @@ def _file_failures( *, evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES, ) -> tuple[str, ...]: - schema_text: Final = _rendered_schema(head_map) - if not schema_text.startswith("{"): - return (schema_text,) backup_failure: Final = ( () if head.backup == head.cost_map else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",) ) - schema_failure: Final = ( - () - if head.schema == schema_text - else ( - f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; " - "run `python ci_cd/generate_model_prices_schema.py` and commit the result", - ) - ) return ( *backup_failure, - *schema_failure, - *( - f"{COST_MAP_PATH} does not validate against its schema: {error}" - for error in validation_errors(head_map, json.loads(schema_text))[:20] - ), *( failure for key, entry in head_map.items() @@ -135,7 +110,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str *( f"bot PRs may only change the cost map files, not {path}" for path in changed_files - if path not in GUARDED_PATHS + if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH ), *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map), *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields), diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py deleted file mode 100644 index e429aae36b9..00000000000 --- a/ci_cd/generate_model_prices_schema.py +++ /dev/null @@ -1,418 +0,0 @@ -from __future__ import annotations - -import json -import sys -from pathlib import Path -from typing import Final, Optional - -import jsonschema - -REPO_ROOT = Path(__file__).parent.parent -PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" -SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" - -SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"}) - -JsonSchema = dict - -NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} -NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} -BOOLEAN: JsonSchema = {"type": "boolean"} -STRING: JsonSchema = {"type": "string"} -TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} -WEEKDAY_PATTERN: Final = ( - r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" -) - -EXTRA_BOOLEAN_KEYS = frozenset( - { - "gemini_native_audio", - "gemini_audio_only_live", - "uses_embed_content", - "use_openai_responses_path", - "bedrock_converse_supports_strict_tools", - "thinking_always_on", - } -) - -HOURS_UTC: Final[JsonSchema] = { - "description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.', - "oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}], -} - -OFF_PEAK_WINDOW: Final[JsonSchema] = { - "type": "object", - "properties": { - "hours_utc": HOURS_UTC, - "weekdays": { - "type": "array", - "description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.", - "items": { - "oneOf": [ - {"type": "integer", "minimum": 1, "maximum": 7}, - {"type": "string", "pattern": WEEKDAY_PATTERN}, - ] - }, - "minItems": 1, - }, - }, - "required": ["hours_utc"], - "additionalProperties": False, -} - -OBJECT_KEYS: dict[str, JsonSchema] = { - "off_peak_pricing": { - "type": "object", - "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.", - "properties": { - "hours_utc": HOURS_UTC, - "windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1}, - "weekday_timezone": { - "type": "string", - "description": "IANA zone the weekdays of each window are read on; defaults to UTC.", - }, - "input_cost_per_token": NONNEG_NUMBER, - "output_cost_per_token": NONNEG_NUMBER, - "output_cost_per_reasoning_token": NONNEG_NUMBER, - "cache_read_input_token_cost": NONNEG_NUMBER, - "cache_creation_input_token_cost": NONNEG_NUMBER, - }, - "anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}], - "additionalProperties": False, - }, - "search_context_cost_per_query": { - "type": "object", - "description": "USD cost per web search query, keyed by search context size.", - "properties": { - "search_context_size_low": NONNEG_NUMBER, - "search_context_size_medium": NONNEG_NUMBER, - "search_context_size_high": NONNEG_NUMBER, - }, - "additionalProperties": False, - }, - "guardrail_cost_per_unit": { - "type": "object", - "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).", - "additionalProperties": NONNEG_NUMBER, - }, - "metadata": { - "type": "object", - "description": "Free-form notes about the entry (e.g. pricing derivation).", - }, - "provider_specific_entry": { - "type": "object", - "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).", - }, -} - -ARRAY_KEYS: dict[str, JsonSchema] = { - "supported_audio_formats": { - "type": "array", - "description": "Audio container formats the model can return.", - "items": {"type": "string", "enum": ["mp3", "wav"]}, - }, - "supported_endpoints": { - "type": "array", - "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", - "items": STRING, - }, - "supported_modalities": { - "type": "array", - "description": "Input modalities the model accepts.", - "items": {"type": "string", "enum": ["text", "image", "audio", "video"]}, - }, - "supported_output_modalities": { - "type": "array", - "description": "Output modalities the model can produce.", - "items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]}, - }, - "reasoning_effort_levels": { - "type": "array", - "description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.", - "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, - }, - "supported_regions": { - "type": "array", - "description": "Cloud regions the model is available in ('global' or region ids).", - "items": STRING, - }, - "tiered_pricing": { - "type": "array", - "description": "Context-length or result-count tiered rates; each tier's costs apply within its range.", - "items": { - "type": "object", - "properties": { - "range": { - "type": "array", - "description": "[min, max] prompt-token span this tier applies to.", - "items": NONNEG_NUMBER, - "minItems": 2, - "maxItems": 2, - }, - "max_results_range": { - "type": "array", - "description": "[min, max] result-count span this tier applies to (search models).", - "items": NONNEG_NUMBER, - "minItems": 2, - "maxItems": 2, - }, - "input_cost_per_token": NONNEG_NUMBER, - "output_cost_per_token": NONNEG_NUMBER, - "output_cost_per_reasoning_token": NONNEG_NUMBER, - "cache_read_input_token_cost": NONNEG_NUMBER, - "cache_creation_input_token_cost": NONNEG_NUMBER, - "input_cost_per_query": NONNEG_NUMBER, - }, - "additionalProperties": False, - }, - }, -} - -INTEGER_KEYS: dict[str, JsonSchema] = { - "max_tokens": { - **NONNEG_INTEGER, - "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.", - }, - "max_input_tokens": { - **NONNEG_INTEGER, - "description": "Maximum prompt/context tokens the model accepts.", - }, - "max_output_tokens": { - **NONNEG_INTEGER, - "description": "Maximum tokens the model can generate in one response.", - }, - "output_vector_size": { - **NONNEG_INTEGER, - "description": "Embedding dimension for embedding models.", - }, - "prompt_cache_min_tokens": { - **NONNEG_INTEGER, - "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", - }, - "token_count_framing": { - **NONNEG_INTEGER, - "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", - }, - "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, - "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, -} - -NUMBER_KEYS: dict[str, JsonSchema] = { - "regional_processing_uplift_multiplier_eu": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).", - }, - "regional_processing_uplift_multiplier_us": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).", - }, - "regional_endpoint_uplift_multiplier": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).", - }, -} - -COST_DESCRIPTIONS: dict[str, str] = { - "input_cost_per_token": "USD per prompt token.", - "output_cost_per_token": "USD per generated token.", - "output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.", - "google_maps_grounding_cost_per_query": ( - "USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit." - ), - "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", - "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", - "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", - "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", -} - - -def cost_description(key: str) -> Optional[str]: - if key in COST_DESCRIPTIONS: - return COST_DESCRIPTIONS[key] - if key.endswith("_flex"): - return "Flex service-tier rate for the same-named base field." - if key.endswith("_priority"): - return "Priority service-tier rate for the same-named base field." - if "_above_" in key: - return "Rate applied once the prompt exceeds the token threshold in the field name." - return None - - -def cost_schema(key: str) -> JsonSchema: - description = cost_description(key) - return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER) - - -def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: - return { - "litellm_provider": { - "type": "string", - "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.", - }, - "mode": { - "type": "string", - "description": "Primary API surface / task type of the model.", - "enum": list(modes), - }, - "source": { - "type": "string", - "description": "URL of the provider pricing/model page this entry was taken from.", - }, - "deprecation_date": { - "type": "string", - "description": "Date the provider deprecates the model, YYYY-MM-DD.", - "format": "date", - "pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$", - }, - "web_search_billing_unit": { - "type": "string", - "description": "Whether web search is billed per query or per prompt.", - "enum": ["per_query", "per_prompt"], - }, - "bedrock_output_config_effort_ceiling": { - "type": "string", - "description": "Highest reasoning effort the Bedrock output_config accepts for this model.", - "enum": ["low", "medium", "high", "max", "xhigh"], - }, - "default_reasoning_effort": { - "type": "string", - "description": ( - "Reasoning effort the provider applies when the request omits reasoning_effort. " - "Gates whether a non-default temperature or the top_p/logprobs sampling params are " - "accepted, which hold only when the effort resolves to 'none'." - ), - "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], - }, - "comment": STRING, - "token_count_fallback_model": { - "type": "string", - "minLength": 1, - "description": "Model that counts tokens for this one when the provider cannot count it.", - }, - "audio_transcription_config": STRING, - "vertex_ai_audio_api": { - "type": "string", - "enum": ["lyria_predict", "lyria_interactions"], - }, - } - - -def classify(key: str, modes: tuple) -> Optional[JsonSchema]: - curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS} - if key in curated: - return curated[key] - if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS: - return BOOLEAN - if "cost" in key: - return cost_schema(key) - return None - - -def build_schema(prices: dict) -> JsonSchema: - entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS} - all_keys = tuple(sorted({key for entry in entries.values() for key in entry})) - modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry})) - unclassified = tuple(key for key in all_keys if classify(key, modes) is None) - if unclassified: - raise SystemExit( - f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. " - f"Add them to the key tables in {Path(__file__).name} and rerun it." - ) - entry_properties = {key: classify(key, modes) for key in all_keys} - return { - "$schema": "https://json-schema.org/draft/2020-12/schema", - "title": "LiteLLM model_prices_and_context_window.json", - "description": ( - "Schema for LiteLLM's model price and context window registry " - "(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). " - "Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, " - "optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. " - "All costs are USD per unit. New optional fields are added regularly, so consumers should " - "ignore unknown fields rather than reject them." - ), - "type": "object", - "properties": { - "sample_spec": { - "type": "object", - "description": ( - "Documentation placeholder illustrating the entry shape; not a real model and not " - "schema-conformant (several values are prose)." - ), - }, - "fallback_generalizations": { - "type": "object", - "description": "Regex rules that generalize unknown model ids to known families; not a model entry.", - "properties": { - "rules": { - "type": "array", - "items": { - "type": "object", - "properties": { - "name": STRING, - "pattern": STRING, - "description": STRING, - }, - "required": ["name", "pattern"], - "additionalProperties": True, - }, - } - }, - "additionalProperties": False, - }, - }, - "additionalProperties": {"$ref": "#/$defs/modelEntry"}, - "$defs": { - "modelEntry": { - "type": "object", - "description": ( - "Pricing, limits, and capability flags for one model. Fields other than litellm_provider " - "are optional; boolean capability flags are simply omitted when unknown or false." - ), - "required": ["litellm_provider"], - "properties": entry_properties, - "additionalProperties": True, - } - }, - } - - -def render(schema: JsonSchema) -> str: - return json.dumps(schema, indent=2) + "\n" - - -def validation_errors(prices: dict, schema: JsonSchema) -> tuple: - validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER) - return tuple( - f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}" - for error in validator.iter_errors(prices) - ) - - -def main() -> int: - check = "--check" in sys.argv[1:] - prices = json.loads(PRICES_PATH.read_text()) - rendered = render(build_schema(prices)) - errors = validation_errors(prices, json.loads(rendered)) - if errors: - print(f"{PRICES_PATH.name} does not validate against the generated schema:") - print("\n".join(errors[:20])) - return 1 - if not check: - SCHEMA_PATH.write_text(rendered) - print(f"wrote {SCHEMA_PATH}") - return 0 - if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered: - print( - f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. " - f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result." - ) - return 1 - print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/gateway/Dockerfile b/gateway/Dockerfile index ee45f293ef8..051b1c30c8f 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -20,6 +20,33 @@ RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${ make -j"$(nproc)" pgbouncer && \ install -m 0755 pgbouncer /usr/local/bin/pgbouncer +# ---------- Rust extension ---------- +# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes +FROM $LITELLM_BUILD_IMAGE AS rust-native + +WORKDIR /app +USER root + +COPY --from=uvbin /uv /uvx /usr/local/bin/ + +RUN for i in 1 2 3; do \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done + +ENV UV_PYTHON_DOWNLOADS=0 + +COPY pyproject.toml README.md rust-toolchain.toml ./ +COPY .cargo ./.cargo +COPY litellm-rust ./litellm-rust + +RUN --mount=type=cache,target=/root/.cache/uv \ + mkdir -p litellm/rust_bridge /wheels && \ + touch litellm/__init__.py litellm/rust_bridge/__init__.py && \ + uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ + python -c "import maturin; maturin.build_wheel('/wheels')" + # ---------- Builder ---------- FROM $LITELLM_BUILD_IMAGE AS builder @@ -66,11 +93,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra bedrock-realtime \ --python python3.13 -# Stage 2 — copy source and install the project + workspace members. +# Stage 2 — copy source and install the workspace members. COPY . . RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-default-groups --no-editable \ + uv sync --frozen --no-install-project --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --group admin-mcp \ @@ -79,6 +106,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra bedrock-realtime \ --python python3.13 +RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \ + uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl + # PYTHONPATH=/app makes the source tree shadow the installed package, so the # compiled Rust extension must live next to the source or it is never imported. RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261015120000_add_environment_variables_revision_trigger/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261015120000_add_environment_variables_revision_trigger/migration.sql new file mode 100644 index 00000000000..67e75896cc4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261015120000_add_environment_variables_revision_trigger/migration.sql @@ -0,0 +1,23 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE SEQUENCE IF NOT EXISTS "litellm_environment_variables_revision_seq"; +GRANT USAGE ON SEQUENCE "litellm_environment_variables_revision_seq" TO PUBLIC; + +CREATE OR REPLACE FUNCTION litellm_bump_environment_variables_revision() RETURNS TRIGGER AS $$ +DECLARE + next_revision BIGINT := nextval(format('%I.litellm_environment_variables_revision_seq', TG_TABLE_SCHEMA)::regclass); +BEGIN + IF TG_OP = 'UPDATE' THEN + NEW."reload_revision" := GREATEST(next_revision, OLD."reload_revision" + 1); + ELSE + NEW."reload_revision" := next_revision; + END IF; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +DROP TRIGGER IF EXISTS "litellm_environment_variables_revision" ON "LiteLLM_Config"; +CREATE TRIGGER "litellm_environment_variables_revision" +BEFORE INSERT OR UPDATE ON "LiteLLM_Config" +FOR EACH ROW WHEN (NEW."param_name" = 'environment_variables') +EXECUTE PROCEDURE litellm_bump_environment_variables_revision(); diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index aa8cb05b170..93bbbb071ba 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -343,6 +343,77 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header( assert_eq!(sent, expected); } +#[rstest] +#[case::anthropic("anthropic", &[AnthropicBeta::DangerousToolUse20260903])] +#[case::deepseek("deepseek", &[])] +#[case::minimax("minimax", &[])] +#[tokio::test] +async fn safeguards_add_the_dangerous_tool_use_beta_only_where_the_beta_map_has_a_column( + call: MessagesCall, + #[case] provider: &str, + #[case] derived: &[AnthropicBeta], +) { + let upstream = upstream([message_response()]).await; + let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]); + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + ..call + }, + json!({"safeguards": safeguards}), + )) + .await; + + let request = only_request(&upstream).await; + let expected: BetaSet = derived + .iter() + .cloned() + .chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) + .collect(); + assert_eq!(sent_betas(&request), expected); + assert_eq!( + request.body_json::().unwrap()["safeguards"], + safeguards + ); +} + +#[rstest] +#[tokio::test] +async fn null_safeguards_are_dropped_instead_of_sent_without_their_beta( + call: MessagesCall, + #[values("anthropic", "azure_ai", "bedrock_mantle", "deepseek", "minimax")] provider: &str, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + ..call + }, + json!({"safeguards": null}), + )) + .await; + + let request = only_request(&upstream).await; + assert_eq!( + sent_betas(&request), + BetaSet::from_iter([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) + ); + assert!( + !request + .body_json::>() + .unwrap() + .contains_key("safeguards") + ); +} + #[rstest] #[tokio::test] async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) { diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index aa926662b7a..bcd1e08ad86 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -109,7 +109,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - update_headers_with_anthropic_beta(headers, request) + merge_beta_headers( + headers, + provider_feature_betas(request, BetaProvider::Anthropic), + ) } } @@ -136,12 +139,20 @@ pub(crate) fn transform_messages_request( messages: strip_encrypted_reasoning_blocks(messages), params: MessagesOptionalParams { context_management, + extra: without_null_safeguards(request.params.extra), ..request.params }, ..request }) } +fn without_null_safeguards(extra: Map) -> Map { + extra + .into_iter() + .filter(|(key, value)| key != "safeguards" || !value.is_null()) + .collect() +} + pub(crate) fn update_headers_with_anthropic_beta( headers: Headers, request: &MessagesRequest, @@ -149,12 +160,34 @@ pub(crate) fn update_headers_with_anthropic_beta( merge_beta_headers(headers, feature_betas(request)) } -/// The betas a request needs on a host other than the first-party API: the Anthropic set -/// filtered and renamed through the host's `BetaProvider` policy. +pub(crate) fn update_headers_with_safeguards_beta( + headers: Headers, + request: &MessagesRequest, + provider: BetaProvider, +) -> Headers { + merge_beta_headers( + headers, + feature_betas(request).union(on_provider(mapped_host_betas(&request.params), provider)), + ) +} + +/// The betas a request needs on a host with a column in the beta map, filtered and renamed +/// through that column. pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet { - feature_betas(request) - .iter() - .filter_map(|beta| beta.on(provider)) + on_provider( + feature_betas(request).union(mapped_host_betas(&request.params)), + provider, + ) +} + +fn on_provider(betas: BetaSet, provider: BetaProvider) -> BetaSet { + betas.iter().filter_map(|beta| beta.on(provider)).collect() +} + +fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet { + carries_safeguards(params) + .then_some(AnthropicBeta::DangerousToolUse20260903) + .into_iter() .collect() } @@ -206,6 +239,13 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool { .is_some_and(|config| config.format.is_some()) } +fn carries_safeguards(params: &MessagesOptionalParams) -> bool { + params + .extra + .get("safeguards") + .is_some_and(|value| !value.is_null()) +} + fn messages_carry_output_config(messages: &[Message]) -> bool { messages .iter() @@ -1025,10 +1065,44 @@ mod tests { json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), &["per-turn-control-2026-07-01"] )] + #[case::safeguards( + json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}), + &[] + )] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } + #[rstest] + #[case::anthropic(BetaProvider::Anthropic, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])] + #[case::vertex_ai(BetaProvider::VertexAi, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])] + #[case::empty_list(BetaProvider::Anthropic, json!([]), &["dangerous-tool-use-2026-09-03"])] + #[case::not_a_list(BetaProvider::Anthropic, json!(1), &["dangerous-tool-use-2026-09-03"])] + #[case::null(BetaProvider::Anthropic, Value::Null, &[])] + #[case::bedrock_converse_has_a_null_cell(BetaProvider::BedrockConverse, json!([{"type": "dangerous_tool_use"}]), &[])] + fn safeguards_add_the_dangerous_tool_use_beta_where_the_beta_map_has_a_cell( + #[case] provider: BetaProvider, + #[case] safeguards: Value, + #[case] expected: &[&str], + ) { + assert_eq!( + provider_feature_betas(&request(json!({"safeguards": safeguards})), provider), + betas(expected) + ); + } + + #[rstest] + fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() { + let client = headers(&[("anthropic-beta", "claude-code-20250219")]); + assert_eq!( + update_headers_with_anthropic_beta( + client.clone(), + &request(json!({"safeguards": [{"type": "dangerous_tool_use"}]})) + ), + client + ); + } + #[rstest] #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])] #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] @@ -1067,6 +1141,11 @@ mod tests { json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), &[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], )] + #[case::safeguards_add_the_dangerous_tool_use_beta_to_the_client_betas( + &[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")], + json!({"safeguards": [{"type": "dangerous_tool_use"}]}), + &[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14")], + )] fn request_headers_merge_the_feature_betas( #[case] input: &[(&str, &str)], #[case] fields: Value, diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index 9eaaa657f6f..3bec142ecf8 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -2,7 +2,7 @@ use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::MessagesRequest, - providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH}, + providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; @@ -12,7 +12,7 @@ use crate::{ common_utils::API_KEY_PLACEMENT, messages::{ handler::shape_anthropic_messages_request, - transformation::{transform_messages_request, update_headers_with_anthropic_beta}, + transformation::{transform_messages_request, update_headers_with_safeguards_beta}, }, }, azure_ai::common_utils::{ @@ -104,7 +104,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - update_headers_with_anthropic_beta(headers, request) + update_headers_with_safeguards_beta(headers, request, BetaProvider::AzureAi) } } @@ -544,6 +544,12 @@ mod tests { &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")] )] #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])] + #[case::safeguards_need_the_dangerous_tool_use_beta( + json!({"safeguards": [{"type": "dangerous_tool_use"}]}), + &[("anthropic-beta", "claude-code-20250219")], + &[("x-api-key", "k"), ("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")] + )] + #[case::null_safeguards_need_no_beta(json!({"safeguards": null}), &[], &[("x-api-key", "k")])] fn request_headers_carry_the_anthropic_feature_betas( #[case] fields: serde_json::Value, #[case] forwarded: &[(&str, &str)], diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 1b99eca975f..e2e64e54664 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -6,7 +6,7 @@ use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; use litellm_llms_types::{ formats::messages::MessagesRequest, - providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH}, + providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; use serde_json::Value; @@ -15,7 +15,7 @@ use crate::{ Error, ErrorDetail, anthropic::messages::{ handler::shape_anthropic_messages_request, - transformation::{transform_messages_request, update_headers_with_anthropic_beta}, + transformation::{transform_messages_request, update_headers_with_safeguards_beta}, }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, @@ -174,7 +174,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()), None => headers, }; - update_headers_with_anthropic_beta(headers, request) + update_headers_with_safeguards_beta(headers, request, BetaProvider::Anthropic) } fn wire_body(&self, body: Value) -> Value { @@ -356,6 +356,21 @@ mod tests { ); } + #[rstest] + #[case::safeguards(json!([{"type": "dangerous_tool_use"}]), Some("dangerous-tool-use-2026-09-03"))] + #[case::null_safeguards(Value::Null, None)] + fn safeguards_need_the_dangerous_tool_use_beta( + #[case] safeguards: Value, + #[case] expected: Option<&str>, + ) { + let request = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": safeguards})).unwrap(); + let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &request); + assert_eq!( + header_value(&headers, "anthropic-beta").map(|value| value.to_string()), + expected.map(str::to_string) + ); + } + #[rstest] #[case::default(json!({}), false)] #[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)] diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 9359184f40e..fefedbbbf5e 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -451,4 +451,27 @@ mod tests { json!({"type": "enabled", "budget_tokens": 1024}) ); } + + #[rstest] + #[case::without_client_betas(&[])] + #[case::with_client_betas(&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")])] + fn safeguards_and_client_betas_reach_deepseek_as_sent(#[case] client: &[(&str, &str)]) { + let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]); + let body = json!({ + "model": "deepseek-v4-pro", + "max_tokens": 100, + "messages": [{"role": "user", "content": "hi"}], + "safeguards": safeguards + }); + let headers: Headers = client + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + assert_eq!( + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG + .request_headers(headers.clone(), &request_from(body.clone())), + headers + ); + assert_eq!(transformed(body)["safeguards"], safeguards); + } } diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index e6a3661a81f..0e665304661 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -213,23 +213,14 @@ fn has_web_search_tool(tools: Option<&[Recognized]>) -> bool { } /// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature -/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled -/// out. +/// betas under Vertex's policy, plus the web search beta Vertex needs spelled out. fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet { let params = &request.params; provider_feature_betas(request, BetaProvider::VertexAi).union( - [ - has_web_search_tool(params.tools.as_deref()) - .then_some(AnthropicBeta::WebSearch20250305), - params - .extra - .get("safeguards") - .is_some_and(|value| !value.is_null()) - .then_some(AnthropicBeta::DangerousToolUse20260903), - ] - .into_iter() - .flatten() - .collect(), + has_web_search_tool(params.tools.as_deref()) + .then_some(AnthropicBeta::WebSearch20250305) + .into_iter() + .collect(), ) } diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index 9fcbcd57a76..0a380788c19 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -2,5 +2,4 @@ For `model_prices_and_context_window.json` validation, we should eventually: -- Remove any schema file like `model_prices_and_context_window.schema.json` - Stop skipping this crate's tests diff --git a/litellm/__init__.py b/litellm/__init__.py index df52808643d..9cc8965b628 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1,34 +1,13 @@ -### Hide pydantic namespace conflict warnings globally ### from __future__ import annotations -import warnings +from litellm._process_defaults import apply_process_defaults as _apply_process_defaults + +_apply_process_defaults() -warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*") -# Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances -# This warning can accumulate during streaming and cause memory leaks -warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*") -# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it -# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked -warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*") ### INIT VARIABLES ######################### import threading import os -# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available -import dotenv as _dotenv - - -def _dev_env_hot_reload_enabled() -> bool: - """The proxy exports this flag when started with ``--reload``. A reloaded - worker is a fresh process that inherits the reloader's environment, so an - edited ``.env`` value stays masked by the stale inherited one unless we - let the file win; overriding makes the edit take effect on reload.""" - return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True" - - -if os.getenv("LITELLM_MODE", "DEV") == "DEV": - _dotenv.load_dotenv(override=_dev_env_hot_reload_enabled()) - from collections.abc import Mapping, Sequence from typing import ( Any, @@ -763,7 +742,7 @@ def is_openai_finetune_model(key: str) -> bool: def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: - for key, value in model_cost_map.items(): + for key, value in tuple(model_cost_map.items()): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): open_ai_chat_completion_models.add(key) elif value.get("litellm_provider") == "text-completion-openai": @@ -1694,6 +1673,9 @@ if TYPE_CHECKING: from .llms.perplexity.decisions.transformation import ( PerplexityDecisionsConfig as PerplexityDecisionsConfig, ) + from .llms.nace.decisions.transformation import ( + NaceDecisionsConfig as NaceDecisionsConfig, + ) from .llms.typesafe.decisions.transformation import ( TypeSafeDecisionsConfig as TypeSafeDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index d47d72ef7a8..8abc24d1417 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -162,6 +162,7 @@ LLM_CONFIG_NAMES: Final = ( "HostedVLLMRerankConfig", "PerplexityDecisionsConfig", "TypeSafeDecisionsConfig", + "NaceDecisionsConfig", "OpenRouterDecisionsConfig", "FireworksAIDecisionsConfig", "CloudflareDecisionsConfig", @@ -729,6 +730,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "HostedVLLMRerankConfig", ), "PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"), + "NaceDecisionsConfig": (".llms.nace.decisions.transformation", "NaceDecisionsConfig"), "TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"), "OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"), "FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"), diff --git a/litellm/_process_defaults.py b/litellm/_process_defaults.py new file mode 100644 index 00000000000..7501215f108 --- /dev/null +++ b/litellm/_process_defaults.py @@ -0,0 +1,24 @@ +import os +import warnings + +import dotenv + + +def _dev_env_hot_reload_enabled() -> bool: + """The proxy exports this flag when started with ``--reload``. A reloaded + worker is a fresh process that inherits the reloader's environment, so an + edited ``.env`` value stays masked by the stale inherited one unless we + let the file win; overriding makes the edit take effect on reload.""" + return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True" + + +def apply_process_defaults() -> None: + warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*") + # Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances + # This warning can accumulate during streaming and cause memory leaks + warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*") + # ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it + # cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked + warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*") + if os.getenv("LITELLM_MODE", "DEV") == "DEV": + dotenv.load_dotenv(override=_dev_env_hot_reload_enabled()) diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 8b7fa1c9347..96c38b69104 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -22,17 +22,28 @@ Environment Variables: - LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional) """ +import asyncio import json import os +import threading +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from functools import cache from importlib.resources import files -from typing import Final +from typing import Final, TypeAlias import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger # Cache for the loaded configuration _BETA_HEADERS_CONFIG: dict | None = None +_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock() +_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config") + +BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] +_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) class GetAnthropicBetaHeadersConfig: @@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict: if _BETA_HEADERS_CONFIG is not None: return _BETA_HEADERS_CONFIG - # Get the URL from environment or use default - from litellm import anthropic_beta_headers_url + with _BETA_HEADERS_CONFIG_LOCK: + if _BETA_HEADERS_CONFIG is not None: + return _BETA_HEADERS_CONFIG - _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) - verbose_logger.debug("Loaded and cached beta headers config") + # Get the URL from environment or use default + from litellm import anthropic_beta_headers_url - return _BETA_HEADERS_CONFIG + _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) + verbose_logger.debug("Loaded and cached beta headers config") + + return _BETA_HEADERS_CONFIG + + +async def aload_beta_headers_config() -> None: + if _BETA_HEADERS_CONFIG is None: + await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config) + + +@cache +def _bundled_beta_headers_config() -> BetaHeadersConfig: + return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()) def reload_beta_headers_config() -> dict: @@ -288,6 +313,14 @@ def is_beta_header_supported( return beta_header in provider_mapping and provider_mapping[beta_header] is not None +def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool: + config: Final = _bundled_beta_headers_config() + aliases: Final = config.get("provider_aliases", {}) + canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider + provider_mapping: Final = config.get(canonical_provider, {}) + return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None + + def get_provider_beta_header( anthropic_beta_header: str, provider: str, diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 11488a03f3f..7089ff78835 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -346,6 +346,7 @@ class LLMCachingHandler: if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): verbose_logger.debug("Checking Async Cache") + logging_obj.reset_caching_details() cached_result = await self._retrieve_from_cache( call_type=call_type, kwargs=kwargs, @@ -357,7 +358,7 @@ class LLMCachingHandler: verbose_logger.debug("Cache Hit!") cache_hit: Final = True end_time: Final = datetime.datetime.now() - model, custom_llm_provider, _, _ = litellm.get_llm_provider( + provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider( model=model, custom_llm_provider=kwargs.get("custom_llm_provider", None), api_base=kwargs.get("api_base", None), @@ -366,7 +367,7 @@ class LLMCachingHandler: cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000 self._update_litellm_logging_obj_environment( logging_obj=logging_obj, - model=model, + model=provider_model, kwargs=kwargs, cached_result=cached_result, is_async=True, @@ -381,7 +382,7 @@ class LLMCachingHandler: call_type=call_type, kwargs=kwargs, logging_obj=logging_obj, - model=model, + model=provider_model, custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) @@ -467,6 +468,7 @@ class LLMCachingHandler: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) print_verbose("Checking Sync Cache") + logging_obj.reset_caching_details() with response_cache_phase("get"): cached_result = litellm.cache.get_cache(**new_kwargs) if is_response_without_output(cached_result): @@ -514,7 +516,7 @@ class LLMCachingHandler: if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks( cached_result=cached_result ): - logging_obj.handle_sync_success_callbacks_for_async_calls( + logging_obj.handle_sync_cache_hit_success_callbacks( result=cached_result, start_time=start_time, end_time=end_time, diff --git a/litellm/constants.py b/litellm/constants.py index 425fc22d340..77f2961d6a1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,4 +1,5 @@ import os +import re import sys from enum import Enum from types import MappingProxyType @@ -18,6 +19,9 @@ MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_PROFILE: Final = "jwt_bearer_obo" MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_SCOPE: Final = "https://graph.microsoft.com/.default" # Replies depend on the caller's delegated identity, which response-cache keys do not include. RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot", "databricks_agent"}) +CLAUDE_CODE_PICKER_PATTERN: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +GATEWAY_CLIENT_HEADER: Final = "x-gateway-client" +CLAUDE_CODE_CLIENT: Final = "claude-code" MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC" GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth" GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type" @@ -391,7 +395,7 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1" NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1" DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3" -BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" +BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime" @@ -1615,6 +1619,7 @@ CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = get_env_int("CLOUDZERO_EXPORT_INTERVA MCP_TOOL_NAME_PREFIX: Final = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = get_env_int("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100) PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096 +INTERNAL_CALLBACK_NAMES: Final = frozenset({"cache", "vector_store_pre_call_hook"}) # Headers to control callbacks X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 5db809e486b..4fcd3bde843 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast from httpx import Response from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -150,6 +150,7 @@ from litellm.utils import ( TranscriptionResponse, cached_get_model_info_helper, token_counter, + unmapped_model_generalization, ) if TYPE_CHECKING: @@ -875,6 +876,16 @@ def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | Non return own_model_info["litellm_provider"] +def _pricing_provider_for_cost_candidate( + candidate_provider: str | None, + request_provider: str | None, + current_provider: object, +) -> str | None: + if candidate_provider != request_provider: + return candidate_provider + return current_provider if isinstance(current_provider, str) else None + + def select_model_name_for_cost_calc( model: str | None, completion_response: object | None, @@ -1438,32 +1449,15 @@ def completion_cost( vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") ) -> float: """ - Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. + Return the USD cost of a LiteLLM call. - Parameters: - completion_response (litellm.ModelResponses): [Required] The response received from a LiteLLM completion request. + Pass `completion_response` to price a response LiteLLM returned. Without one, pass `model` plus + `messages` / `prompt` and `completion`, and token counts are computed from the text. Prices come from + `litellm.model_cost` unless `custom_cost_per_token` or `custom_cost_per_second` is given. `total_time` + is in milliseconds and only matters for per-second pricing such as Replicate and SageMaker. - [OPTIONAL PARAMS] - model (str): Optional. The name of the language model used in the completion calls - prompt (str): Optional. The input prompt passed to the llm - completion (str): Optional. The output completion text from the llm - total_time (float, int): Optional. (Only used for Replicate LLMs) The total time used for the request in seconds - custom_cost_per_token: Optional[CostPerToken]: the cost per input + output token for the llm api call. - custom_cost_per_second: Optional[float]: the cost per second for the llm api call. - - Returns: - float: The cost in USD dollars for the completion based on the provided parameters. - - Exceptions: - Raises exception if model not in the litellm model cost map. Register model, via custom pricing or PR - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json - - - Note: - - If completion_response is provided, the function extracts token information and the model name from it. - - If completion_response is not provided, the function calculates token counts based on the model and input text. - - The cost is calculated based on the model, prompt tokens, and completion tokens. - - For certain models containing "togethercomputer" in the name, prices are based on the model size. - - For un-mapped Replicate models, the cost is calculated based on the total time used for the request. + Raises if no candidate model name resolves to pricing. Register the model via `litellm.register_model` + or add it to model_prices_and_context_window.json. """ try: call_type = infer_call_type(call_type, completion_response) or "completion" @@ -1561,14 +1555,20 @@ def completion_cost( ) requested_model: Final = model - potential_model_names: Final = [ - selected_model, - _get_response_model(completion_response), - ] - if model is not None: - potential_model_names.append(model) + request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + raw_candidates: Final = ( + ( + selected_model, + base_model_provider if base_model_provider is not None else request_pricing_provider, + ), + (_get_response_model(completion_response), request_pricing_provider), + *(((model, request_pricing_provider),) if model is not None else ()), + ) + potential_model_names: Final = ( + raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates) + ) - for idx, model in enumerate(potential_model_names): + for idx, (model, pricing_provider) in enumerate(potential_model_names): try: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("selected model name for cost calculation: %s", model) @@ -1913,13 +1913,6 @@ def completion_cost( if litellm_logging_obj is not None: request_model_for_cost = litellm_logging_obj.model - pricing_provider: str | None = ( - base_model_provider - if base_model_provider is not None and idx == 0 - else custom_llm_provider - if isinstance(custom_llm_provider, str) - else None - ) ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -1927,7 +1920,11 @@ def completion_cost( model=model, prompt_tokens=prompt_tokens or 0, completion_tokens=completion_tokens or 0, - custom_llm_provider=pricing_provider, + custom_llm_provider=_pricing_provider_for_cost_candidate( + candidate_provider=pricing_provider, + request_provider=request_pricing_provider, + current_provider=custom_llm_provider, + ), response_time_ms=response_time_ms, region_name=None if explicit_pricing else region_name, custom_cost_per_second=custom_cost_per_second, @@ -2026,7 +2023,11 @@ def completion_cost( if cost_per_token_usage_object is not None and model: _token_type_breakdown = get_token_type_cost_breakdown( model=model, - custom_llm_provider=pricing_provider, + custom_llm_provider=_pricing_provider_for_cost_candidate( + candidate_provider=pricing_provider, + request_provider=request_pricing_provider, + current_provider=custom_llm_provider, + ), usage=cost_per_token_usage_object, service_tier=service_tier, data_residency=data_residency, @@ -2066,7 +2067,7 @@ def completion_cost( ) if idx == len(potential_model_names) - 1: raise e - raise Exception(f"Unable to calculat cost for received potential model names - {potential_model_names}") + raise Exception(f"Unable to calculate cost for received potential model names - {potential_model_names}") except Exception as e: raise e @@ -2081,13 +2082,20 @@ def get_response_cost_from_hidden_params( additional_headers: Final = _hidden_params_dict.get("additional_headers", {}) if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers: - response_cost: Final = additional_headers["llm_provider-x-litellm-response-cost"] - if response_cost is None: - return None - return float(additional_headers["llm_provider-x-litellm-response-cost"]) + return _billable_header_cost(additional_headers["llm_provider-x-litellm-response-cost"]) return None +def _billable_header_cost(value: object) -> float | None: + if not isinstance(value, (str, int, float)): + return None + try: + cost: Final = float(value) + except ValueError: + return None + return cost if math.isfinite(cost) and cost >= 0 else None + + def response_cost_calculator( response_object: ModelResponse | EmbeddingResponse @@ -2281,6 +2289,25 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn return None +_CostCandidate: TypeAlias = tuple[str | None, str | None] + + +def _is_unpriced_generalized_candidate(candidate: _CostCandidate) -> bool: + model, pricing_provider = candidate + if model is None: + return False + rule_info: Final = unmapped_model_generalization(model, pricing_provider) + return rule_info is not None and not _cost_map_entry_prices_anything(rule_info) + + +def _priced_candidates_before_unpriced_rules(candidates: tuple[_CostCandidate, ...]) -> tuple[_CostCandidate, ...]: + unpriced: Final = tuple(_is_unpriced_generalized_candidate(candidate) for candidate in candidates) + return ( + *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if not is_unpriced), + *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if is_unpriced), + ) + + def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None: raw_entry: Final = litellm.model_cost.get(key) return raw_entry if isinstance(raw_entry, Mapping) else None diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 040d80e1ac0..4943b22e1fb 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -57,7 +57,7 @@ from litellm.types.decisions import ( UnsupportedDecisionsRequest, ) from litellm.types.utils import LlmProviders, ProviderSpecificHeader -from litellm.utils import ProviderConfigManager +from litellm.utils import ProviderConfigManager, get_model_cost_key SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]] OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]] @@ -667,7 +667,10 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci { "model": f"{call.custom_llm_provider}/{call.model}", "custom_llm_provider": call.custom_llm_provider, - "provider_response_model": f"{call.custom_llm_provider}/{call.model}", + "provider_response_model": ( + get_model_cost_key(f"{call.custom_llm_provider}/{response.model}") + or f"{call.custom_llm_provider}/{call.model}" + ), **({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}), } ) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e2057e4b627..ed6d89397a0 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -498,6 +498,7 @@ class MCPClient: elicitation_callback: "ElicitationCallback | None" = None, logging_callback: Callable | None = None, protocol_version: MCPUpstreamProtocol = "auto", + upstream_transport: httpx2.AsyncBaseTransport | None = None, ): self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol]( MCPUpstreamProtocol @@ -524,6 +525,7 @@ class MCPClient: self._sampling_callback: Callable | None = sampling_callback self._elicitation_callback: ElicitationCallback | None = elicitation_callback self._logging_callback: Callable | None = logging_callback + self._upstream_transport: httpx2.AsyncBaseTransport | None = upstream_transport self._session_callbacks_override: _SessionCallbackMap | None = None # handle the basic auth value if provided if auth_value: @@ -610,7 +612,11 @@ class MCPClient: return stdio_client(server_params), None if self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() - httpx_client_factory = self._create_httpx_client_factory() + httpx_client_factory = ( + self._create_httpx_client_factory() + if self._upstream_transport is None + else self._create_httpx_client_factory(transport=self._upstream_transport) + ) return ( sse_client( url=self.server_url, @@ -622,7 +628,11 @@ class MCPClient: ) # HTTP transport (default) headers = self._get_auth_headers() - httpx_client_factory = self._create_httpx_client_factory() + httpx_client_factory = ( + self._create_httpx_client_factory() + if self._upstream_transport is None + else self._create_httpx_client_factory(transport=self._upstream_transport) + ) verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) http_client = httpx_client_factory( headers=headers, diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 90383ddb60d..fb0723289ee 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -6,9 +6,10 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, cast from pydantic import ConfigDict, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -106,6 +107,14 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i # Cap on credential-scoped providers held at once; each one owns an exporter thread. _MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256 + +class _CacheHitView(TypedDict, total=False): + cache_hit: ReadOnly[object] + saved_cache_cost: ReadOnly[object] + + +_CACHE_HIT_VIEW: Final = TypeAdapter(_CacheHitView) + # Dedicated so a slow exporter shutdown cannot starve the shared logging executor. _PROVIDER_SHUTDOWN_EXECUTOR: Final = ThreadPoolExecutor(max_workers=4, thread_name_prefix="OtelProviderShutdown") @@ -2434,6 +2443,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): litellm_params=litellm_params, ) # Cost breakdown tracking + cache_view: Final = _CACHE_HIT_VIEW.validate_python(standard_logging_payload) + cache_hit: Final = cache_view.get("cache_hit") is True cost_breakdown: Final[CostBreakdown | None] = standard_logging_payload.get("cost_breakdown") if cost_breakdown: for key, value in cost_breakdown.items(): @@ -2443,6 +2454,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): key=f"gen_ai.cost.{key}", value=value, ) + self.safe_set_attribute(span=span, key=LiteLLM.CACHE_HIT, value=cache_hit) + if cache_hit: + saved_cache_cost: Final = cache_view.get("saved_cache_cost") + if isinstance(saved_cache_cost, (int, float)): + self.safe_set_attribute(span=span, key=LiteLLM.SAVED_CACHE_COST, value=saved_cache_cost) ############################################# ########## LLM Request Attributes ########### ############################################# diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index a4cb0a84541..8779153959f 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -100,23 +100,19 @@ def _carried_keys( return frozenset(key for key, _value in attributes) -def attribute_budget(span: Span, reserved: int, overwrites: frozenset[str] = frozenset()) -> int | None: - """How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more. +def carried_keys(span: Span) -> frozenset[str]: + """The attribute keys ``span`` already holds, empty for a span that exposes none.""" + if not isinstance(span, ReadableSpan): + return frozenset() + return _carried_keys(span.attributes or ()) - ``overwrites`` are the mapped keys already present on ``span``: setting one - replaces the value in place and consumes no slot against the limit, so only - the genuinely new pre-existing keys reduce the budget. Counting the - overwritten ones too reserves slots the fit can never spend and sheds - indexed message attributes for nothing. - """ + +def attribute_budget(span: Span, reserved: int) -> int | None: + """How many new attribute keys fit on ``span`` next to what it already carries and ``reserved`` more.""" limit: Final = span_attribute_limit(span) if limit is None: return None - if not isinstance(span, ReadableSpan): - return limit - reserved - existing_keys: Final = _carried_keys(span.attributes or ()) - fresh: Final = len(existing_keys - overwrites) if overwrites else len(existing_keys) - return limit - fresh - reserved + return limit - len(carried_keys(span)) - reserved def stamp_error( @@ -303,8 +299,14 @@ class SpanEmitter: ) stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES reserved: Final = len(stamped_later.keys() - mapped.keys()) - for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved, frozenset(mapped))).items(): - span.set_attribute(key, value) + overwrites: Final = carried_keys(span) & mapped.keys() + fitted: Final = fit_indexed_messages( + MappingProxyType({key: value for key, value in mapped.items() if key not in overwrites}), + attribute_budget(span, reserved), + ) + for key, value in mapped.items(): + if key in overwrites or key in fitted: + span.set_attribute(key, value) if error: stamped: Final = stamp_error(span, error) if stamped is not None and self._event_recorder is not None and role is SpanRole.LLM_CALL: diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 9eb9cb40caf..73448221acb 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -538,8 +538,12 @@ class OpenTelemetryV2(CustomLogger): # logger is a success/failure callback only, so ``pre_call`` never reaches it # and no carrier exists. The payload plus the request-level provider-handoff # stamp (``upstream_started``) is the affirmative signal of a real call; a - # gate rejection carries ``is_no_upstream_call`` and gets no span. - if carrier is None and (call.is_no_upstream_call or not call.upstream_started or call.payload is None): + # gate rejection carries ``is_no_upstream_call`` and gets no span. A cache + # hit never hands off to a provider, so ``cache_hit`` substitutes for the + # upstream stamp, since the payload proves the call completed. + if carrier is None and ( + call.is_no_upstream_call or call.payload is None or (not call.upstream_started and not call.cache_hit) + ): return None try: return self._finish_carrier(carrier, call, start_time, end_time) diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 0ac8b47e239..4e07d090edd 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -91,6 +91,8 @@ class GenAIMapper: f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount, f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent, f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount, + LiteLLM.CACHE_HIT: lambda d: d.cache_hit, + LiteLLM.SAVED_CACHE_COST: lambda d: d.saved_cache_cost, LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming, LiteLLM.REQUEST_ROUTE: lambda d: d.request_route, LiteLLM.REQUEST_PURPOSE: lambda d: d.request_purpose, diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index cbbb1a9f6d4..a3151c24519 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -249,6 +249,10 @@ class LLMCallEvent: # ``InternalCallOrigin`` such as ``autorouter_classifier``), ``None`` for the # caller's own provider attempt. purpose: str | None + # True for a response served from the litellm response cache. A cache hit + # never hands off to a provider, so ``upstream_started`` stays False even + # though the call completed and is worth a span. + cache_hit: bool # A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The # span is renamed from the typed payload at close (``finish_span``); this only # needs to be reasonable for a span that never gets closed (a leak). @@ -273,6 +277,7 @@ class LLMCallEvent: upstream_started=kwargs.get("api_call_start_time") is not None, upstream_start_seconds=_epoch_seconds(kwargs.get("api_call_start_time")), purpose=internal_call_origin(payload, kwargs), + cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), trace=trace, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 101e6285946..5bca8e01aec 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -441,6 +441,8 @@ class LLMCallSpanData: embedding_output: EmbeddingOutput | None = None promoted_metadata: Mapping[str, str] = field(default_factory=lambda: _EMPTY_METADATA) routing_attributes: Mapping[str, RoutingAttributeValue] = field(default_factory=lambda: MappingProxyType({})) + cache_hit: bool | None = None + saved_cache_cost: float | None = None @classmethod def from_standard_logging_payload( @@ -456,6 +458,7 @@ class LLMCallSpanData: metadata_keys: tuple[str, ...] = (), ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) + cache_hit: Final = as_bool(payload.get("cache_hit")) is True # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown # The single parse of the request's metadata — the request-vs-provider # model split, the response model, api base, and identity all come from # here rather than being re-derived from the raw payload dicts. @@ -490,7 +493,7 @@ class LLMCallSpanData: error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))), - server=ServerInfo.from_api_base(context.api_base), + server=None if cache_hit else ServerInfo.from_api_base(context.api_base), identity=context.identity, promoted_metadata=allowlisted_metadata(context.identity.metadata, metadata_keys), is_streaming=as_bool(payload.get("stream")), @@ -507,6 +510,8 @@ class LLMCallSpanData: session_id=session_id or None, embedding_output=embedding_output if capture_content else None, routing_attributes=routing_decision_attributes(decision), + cache_hit=cache_hit, + saved_cache_cost=as_float(payload.get("saved_cache_cost")) if cache_hit else None, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown ) diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 38cf7266593..b7f955a75ed 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -340,6 +340,8 @@ class LiteLLM: # semconv key for an MCP server's *name* (the convention uses ``server.address`` # for its network location), so it lives under the vendor namespace. MCP_SERVER_NAME: Final = "litellm.mcp.server.name" + CACHE_HIT: Final = "litellm.cache_hit" + SAVED_CACHE_COST: Final = "litellm.saved_cache_cost" class Metric: diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 97a1073d7c0..43250fd0fd4 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -259,6 +259,7 @@ _POSITIVE_SERIES_CAP: Final[TypeAdapter[int]] = TypeAdapter(Annotated[int, Field _POSITIVE_SERIES_TTL: Final[TypeAdapter[float]] = TypeAdapter(Annotated[float, Field(gt=0)]) _SERIES_CLEANUP_INTERVAL: Final[TypeAdapter[float]] = TypeAdapter(Annotated[float, Field(ge=0)]) _DEFAULT_SERIES_CLEANUP_INTERVAL_SECONDS: Final = 60.0 +_ADDITIONAL_HEADERS: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) def _number_or_none(value: object, limit: TypeAdapter[_SeriesLimitT]) -> _SeriesLimitT | None: @@ -270,6 +271,15 @@ def _number_or_none(value: object, limit: TypeAdapter[_SeriesLimitT]) -> _Series return None +def _gauge_reading_or_none(value: object) -> float | None: + if not isinstance(value, str | int | float): + return None + try: + return float(value) + except (ValueError, OverflowError): + return None + + def _positive_or_ignored(setting: str, value: object, limit: TypeAdapter[_SeriesLimitT]) -> _SeriesLimitT | None: if value is None: return None @@ -3350,12 +3360,11 @@ class PrometheusLogger(CustomLogger): model_group=enum_values.model_group, ) - remaining_requests: int | None = None - remaining_tokens: int | None = None - if additional_headers := standard_logging_payload["hidden_params"]["additional_headers"]: - # OpenAI / OpenAI Compatible headers - remaining_requests = additional_headers.get("x_ratelimit_remaining_requests", None) - remaining_tokens = additional_headers.get("x_ratelimit_remaining_tokens", None) + additional_headers: Final = _ADDITIONAL_HEADERS.validate_python( + standard_logging_payload["hidden_params"]["additional_headers"] or {} + ) + remaining_requests: Final = _gauge_reading_or_none(additional_headers.get("x_ratelimit_remaining_requests")) + remaining_tokens: Final = _gauge_reading_or_none(additional_headers.get("x_ratelimit_remaining_tokens")) if litellm_overhead_time_ms := standard_logging_payload["hidden_params"].get("litellm_overhead_time_ms"): _labels = prometheus_label_factory( @@ -3373,7 +3382,7 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) - if remaining_requests: + if remaining_requests is not None: """ "model_group", "api_provider", @@ -3387,7 +3396,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) - if remaining_tokens: + if remaining_tokens is not None: _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_tokens_metric"), enum_values=enum_values, diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index ab386e10951..ffe2c4433a0 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -44,7 +44,7 @@ from litellm.litellm_core_utils.private_json import ( stage_private_json, write_private_json, ) -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel @dataclass(frozen=True, slots=True) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cc1f6fef019..e706bcadef8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -318,7 +318,7 @@ def remove_items_at_indices(items: _PoppableByIndex | None, indices: Iterable[in items.pop(index) -def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: dict) -> dict: +def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: Mapping[str, object]) -> dict: """ Helper to get litellm metadata for spend tracking @@ -326,7 +326,7 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad and user_api_key values are in 'metadata'. """ potential_spend_tracking_metadata_substring: Final = "user_api_key" - for key, value in metadata.items(): + for key, value in dict(metadata).items(): if potential_spend_tracking_metadata_substring in key: litellm_metadata[key] = value return litellm_metadata diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index f8aadc4258d..c52bc534058 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1196,6 +1196,9 @@ def _map_sagemaker_exception( ) +_VERTEX_403_STATUS_CODE_PATTERN: Final = re.compile(r"(? None: + status_code: Final[object] = getattr(original_exception, "status_code", None) + body_code: Final[int | None] = _get_body_error_code(error_str) if "Vertex AI API has not been used in project" in error_str or "Unable to find your project" in error_str: raise BadRequestError( message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}", @@ -1252,7 +1257,11 @@ def _map_vertex_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, ) - elif "403" in error_str: + elif ( + (not isinstance(status_code, int) or status_code in (403, 500)) + and body_code in (None, 403) + and _VERTEX_403_STATUS_CODE_PATTERN.search(error_str) + ): raise BadRequestError( message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}", model=model, @@ -1306,7 +1315,7 @@ def _map_vertex_exception( elif ( isinstance(getattr(original_exception, "status_code", None), int) and 500 <= original_exception.status_code < 600 - and _get_body_error_code(error_str) == 429 + and body_code == 429 ): # upstream gateway wraps a 429 inside a 5xx envelope # e.g. HTTP 500/503 with {"error":{"code":429,...}}. @@ -2413,6 +2422,35 @@ def _map_exception_by_status( ) +_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) +_URL_SCHEME: Final = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*") +_HOST_WITH_OPTIONAL_PORT: Final = re.compile( + r"(?:(?:[A-Za-z0-9-]+\.)+[A-Za-z0-9-]+|localhost|\[[0-9A-Fa-f:.]+\])(?::\d+)?" +) + + +def with_deployment_api_base(litellm_params: object) -> Mapping[str, object]: + params: Final = as_str_mapping(litellm_params) or _NO_LITELLM_PARAMS + metadata: Final = as_str_mapping(params.get("metadata")) + deployment_api_base: Final = None if metadata is None else metadata.get("api_base") + if not isinstance(deployment_api_base, str) or not deployment_api_base: + return params + return MappingProxyType({**params, "api_base": deployment_api_base}) + + +def _api_base_without_credentials(api_base: str) -> str: + scheme, separator, _ = api_base.partition("://") + scheme_prefix: Final = scheme + separator if separator and _URL_SCHEME.fullmatch(scheme) else "" + authority_and_rest: Final = api_base[len(scheme_prefix) :] + authority: Final = re.split(r"[/?#]", authority_and_rest, maxsplit=1)[0] + after_authority: Final = authority_and_rest[len(authority) :] + host: Final = authority.rpartition("@")[2] + if "@" in after_authority and _HOST_WITH_OPTIONAL_PORT.fullmatch(host) is None: + return "" + path: Final = re.split(r"[?#]", after_authority, maxsplit=1)[0] + return f"{scheme_prefix}{host}{path}" + + def _is_guardrail_block(original_exception: Exception) -> bool: from litellm.integrations.custom_guardrail import is_guardrail_intervention @@ -2483,8 +2521,9 @@ def exception_type( if isinstance(custom_llm_provider, str) and len(custom_llm_provider) > 0: exception_provider = custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" - if _api_base: - extra_information += f"\nAPI Base: `{_api_base}`" + shown_api_base: Final = _api_base_without_credentials(_api_base) if _api_base else "" + if shown_api_base: + extra_information += f"\nAPI Base: `{shown_api_base}`" if messages and len(messages) > 0 and litellm.redact_messages_in_exceptions is False: extra_information += f"\nMessages: `{messages}`" diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 6e3769260cc..073e3a3455e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -324,10 +324,16 @@ def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _Fetc return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag) -def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: +def _split_url(url: str) -> urllib.parse.SplitResult | ModelCostMapReloadUnavailable: + try: + return urllib.parse.urlsplit(url) + except ValueError as error: + return ModelCostMapReloadUnavailable(reason=f"invalid URL {url}: {error}") + + +def _attempt_fetch_s3(parts: urllib.parse.SplitResult, url: str, timeout: int) -> _FetchAttemptOutcome: from botocore.exceptions import BotoCoreError, ClientError - parts: Final = urllib.parse.urlsplit(url) bucket: Final = parts.netloc key: Final = parts.path.lstrip("/") if not bucket or not key: @@ -343,9 +349,11 @@ def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = urllib.parse.urlsplit(url) + parts: Final = _split_url(url) + if isinstance(parts, ModelCostMapReloadUnavailable): + return parts if parts.scheme == "s3": - return await asyncio.to_thread(_attempt_fetch_s3, url, timeout) + return await asyncio.to_thread(_attempt_fetch_s3, parts, url, timeout) try: response: Final = await client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: @@ -354,9 +362,11 @@ async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _Fe def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = urllib.parse.urlsplit(url) + parts: Final = _split_url(url) + if isinstance(parts, ModelCostMapReloadUnavailable): + return parts if parts.scheme == "s3": - return _attempt_fetch_s3(url, timeout) + return _attempt_fetch_s3(parts, url, timeout) try: response: Final = client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: @@ -547,22 +557,39 @@ def get_model_cost_map_loaded_at() -> "datetime | None": return _cost_map_source_info.loaded_at -def _expand_model_aliases(model_cost: dict) -> dict: +class AliasedCostEntry(dict[str, object]): + __slots__ = ("names",) + + def __init__(self, entry: Mapping[str, object], names: tuple[str, ...]) -> None: + super().__init__(entry) + self.names: tuple[str, ...] = names + + +def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]: + return MappingProxyType( + { + owner: (owner, *(alias for alias, alias_owner in alias_owners.items() if alias_owner == owner)) + for owner in dict.fromkeys(alias_owners.values()) + } + ) + + +def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]: """ Expand ``aliases`` lists in model cost entries into top-level entries. - Each alias gets a reference to the **same** dict object as the canonical - entry (zero memory overhead). The ``aliases`` key is removed from the - entry so downstream code never sees it. + Each alias gets a reference to the **same** ``AliasedCostEntry`` as the + canonical entry, which carries every name it is served under. The + ``aliases`` key is removed from the entry so downstream code never sees it. If an alias collides with an existing canonical entry the alias is skipped and a warning is logged. """ - aliases_to_add: Final[dict[str, dict]] = {} + alias_owners: Final[dict[str, str]] = {} keys_with_aliases: Final[list[str]] = [] for model_name, model_info in model_cost.items(): - aliases: list | None = model_info.get("aliases") + aliases: object = model_info.get("aliases") if aliases is None: continue keys_with_aliases.append(model_name) @@ -584,7 +611,7 @@ def _expand_model_aliases(model_cost: dict) -> dict: model_name, ) continue - if alias in aliases_to_add: + if alias in alias_owners: verbose_logger.warning( "LiteLLM model alias conflict: alias '%s' (from '%s') " "was already claimed by another entry — skipping.", @@ -592,13 +619,14 @@ def _expand_model_aliases(model_cost: dict) -> dict: model_name, ) continue - aliases_to_add[alias] = model_info # same dict reference + alias_owners[alias] = model_name # Remove the ``aliases`` key from entries so it doesn't pollute model info for key in keys_with_aliases: model_cost[key].pop("aliases", None) - model_cost.update(aliases_to_add) + for owner, names in _alias_groups(alias_owners).items(): + model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names))) return model_cost diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8c310011556..ea003034777 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -13,6 +13,7 @@ import sys import time import traceback from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType @@ -41,6 +42,7 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, + INTERNAL_CALLBACK_NAMES, PROVIDER_REQUEST_ID_HEADERS, REDACTED_BY_LITELLM, ) @@ -557,6 +559,12 @@ def mask_api_base_credentials(api_base: str) -> str: return api_base[:key_end] + "*" * 5 + api_base[-4:] +@dataclass(frozen=True, slots=True) +class UncachedPrice: + cost: float | None + cost_breakdown: CostBreakdown | None + + def _timestamp_seconds(moment: object) -> float | None: if isinstance(moment, datetime.datetime): return moment.timestamp() @@ -810,6 +818,10 @@ class Logging(LiteLLMLoggingBaseClass): new_call.sync_streaming_chunks = [] return new_call + def reset_caching_details(self) -> None: + """Forget an earlier cache lookup's result, so a call that looks up again reports only its latest.""" + self.caching_details = None + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) @@ -1899,6 +1911,7 @@ class Logging(LiteLLMLoggingBaseClass): cache_hit: bool | None = None, litellm_model_name: str | None = None, router_model_id: str | None = None, + ignore_stamped_cost: bool = False, ) -> float | None: """ Calculate response cost using result + logging object variables. @@ -1934,7 +1947,9 @@ class Logging(LiteLLMLoggingBaseClass): ): hidden_params: Final = result_hidden_params if ( - "response_cost" in hidden_params and hidden_params["response_cost"] is not None + not ignore_stamped_cost + and "response_cost" in hidden_params + and hidden_params["response_cost"] is not None ): # use cost if already calculated self._record_zero_cost_diagnostic( priced_result, @@ -2060,6 +2075,20 @@ class Logging(LiteLLMLoggingBaseClass): _response_cost_calculator = response_cost_calculator + def uncached_price(self, result: object) -> UncachedPrice: + previous_cost_breakdown: Final = self.cost_breakdown + previous_billed_token_rates: Final = self.billed_token_rates + debug_key: Final = "response_cost_failure_debug_information" + had_failure_debug_info: Final = debug_key in self.model_call_details + try: + cost: Final = self.response_cost_calculator(result=result, cache_hit=False, ignore_stamped_cost=True) + return UncachedPrice(cost=cost, cost_breakdown=self.cost_breakdown) + finally: + self.cost_breakdown = previous_cost_breakdown + self.billed_token_rates = previous_billed_token_rates + if not had_failure_debug_info: + self.model_call_details.pop(debug_key, None) + def _record_zero_cost_diagnostic( self, result: object, @@ -2303,6 +2332,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, cache_hit=cache_hit, + for_async_request=True, **kwargs, ) @@ -2545,7 +2575,8 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, build_logging_payload: bool = True, - ): + publish_logging_payload: bool = True, + ) -> StandardLoggingPayload | None: """Resolve hidden params, compute response cost, and emit the standard logging payload.""" self._surface_response_headers_from_result(logging_result) hidden_params: Final = getattr(logging_result, HIDDEN_PARAMS_ATTR, {}) @@ -2572,17 +2603,17 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.response_cost_calculator(result=logging_result) if not build_logging_payload: - return + return None - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - logging_result, start_time, end_time - ) - - standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( - "standard_logging_object" + standard_logging_payload: Final = self._build_standard_logging_payload( + logging_result, start_time, end_time, isolate_call_details=not publish_logging_payload ) + if not publish_logging_payload: + return standard_logging_payload + self.model_call_details["standard_logging_object"] = standard_logging_payload if standard_logging_payload is not None: emit_standard_logging_payload(standard_logging_payload) + return standard_logging_payload async def _prepare_baseline_cache_estimate(self, response_obj: object) -> None: if self.baseline_cache_context is None: @@ -2598,12 +2629,16 @@ class Logging(LiteLLMLoggingBaseClass): await invalidate_baseline_cache(self, reason, completed=completed) def _build_standard_logging_payload( - self, init_response_obj: object, start_time: dt_object, end_time: dt_object + self, + init_response_obj: object, + start_time: dt_object, + end_time: dt_object, + isolate_call_details: bool = False, ) -> StandardLoggingPayload | None: """Build StandardLoggingPayload and accumulate its construction time.""" _start: Final = time.time() payload: Final = get_standard_logging_object_payload( - kwargs=self.model_call_details, + kwargs=dict(self.model_call_details) if isolate_call_details else self.model_call_details, init_response_obj=init_response_obj, start_time=start_time, end_time=end_time, @@ -2620,11 +2655,16 @@ class Logging(LiteLLMLoggingBaseClass): emit_standard_logging_payload(stored_payload) def _transform_usage_objects(self, result: object) -> object: + return self._transform_usage_objects_for_payload(result, self.model_call_details.get("standard_logging_object")) + + def _transform_usage_objects_for_payload( + self, result: object, standard_logging_payload: StandardLoggingPayload | None + ) -> object: if isinstance(result, ResponsesAPIResponse): result = result.model_copy() transformed_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(result.usage) setattr(result, "usage", transformed_usage) - if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: + if standard_logging_payload is not None: response_dict: Final = result.model_dump() if hasattr(result, "model_dump") else dict(result) # Ensure usage is properly included with transformed chat format if transformed_usage is not None: @@ -2633,7 +2673,7 @@ class Logging(LiteLLMLoggingBaseClass): if hasattr(transformed_usage, "model_dump") else dict(transformed_usage) ) - standard_logging_payload["response"] = response_dict + standard_logging_payload["response"] = response_dict # rebind-ok: mirrors the transformed usage elif isinstance(result, TranscriptionResponse): from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( TranscriptionUsageObjectTransformation, @@ -2673,6 +2713,27 @@ class Logging(LiteLLMLoggingBaseClass): standard_logging_object: StandardLoggingPayload | None = None, build_logging_payload: bool = True, ) -> tuple[datetime.datetime, datetime.datetime, object]: + prepared: Final = self._prepare_success_logging( + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + standard_logging_object=standard_logging_object, + build_logging_payload=build_logging_payload, + publish_logging_payload=True, + ) + return prepared[0], prepared[1], prepared[2] + + def _prepare_success_logging( + self, + result: object, + start_time: datetime.datetime | None, + end_time: datetime.datetime | None, + cache_hit: bool | None, + standard_logging_object: StandardLoggingPayload | None, + build_logging_payload: bool, + publish_logging_payload: bool, + ) -> tuple[datetime.datetime, datetime.datetime, object, StandardLoggingPayload | None]: try: if start_time is None: start_time = self.start_time @@ -2705,29 +2766,21 @@ class Logging(LiteLLMLoggingBaseClass): ): result = logging_result - if standard_logging_object is None and result is not None and self.stream is not True: - if self._is_recognized_call_type_for_logging(logging_result=logging_result) or isinstance( - logging_result, (dict, list) - ): - self._process_hidden_params_and_response_cost( - logging_result=logging_result, - start_time=start_time, - end_time=end_time, - build_logging_payload=build_logging_payload, - ) - elif standard_logging_object is not None: - self.model_call_details["standard_logging_object"] = standard_logging_object - else: - # Streaming reaches here before its cost is known, so the cost - # is seeded to None, but only when nothing has already - # established one. A stream that assembles into a response - # object recomputes the cost right after this; a pass-through - # stream cannot (its body is opaque) and carries the cost its - # upstream reported in the response headers, which an - # unconditional reset would discard. - self.model_call_details.setdefault("response_cost", None) + standard_logging_payload: Final = self._resolve_success_logging_payload( + result=result, + logging_result=logging_result, + start_time=start_time, + end_time=end_time, + standard_logging_object=standard_logging_object, + build_logging_payload=build_logging_payload, + publish_logging_payload=publish_logging_payload, + ) - result = self._transform_usage_objects(result=result) + result = ( + self._transform_usage_objects(result=result) + if publish_logging_payload + else self._transform_usage_objects_for_payload(result, standard_logging_payload) + ) if ( litellm.max_budget @@ -2746,10 +2799,47 @@ class Logging(LiteLLMLoggingBaseClass): standard_built_in_tools_params=self.standard_built_in_tools_params, ) - return start_time, end_time, result + return start_time, end_time, result, standard_logging_payload except Exception as e: raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {e}") + def _resolve_success_logging_payload( + self, + result: object, + logging_result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + standard_logging_object: StandardLoggingPayload | None, + build_logging_payload: bool, + publish_logging_payload: bool, + ) -> StandardLoggingPayload | None: + if standard_logging_object is not None: + if not publish_logging_payload: + return standard_logging_object.copy() + self.model_call_details["standard_logging_object"] = standard_logging_object + return standard_logging_object + if result is None or self.stream is True: + # Streaming reaches here before its cost is known, so the cost + # is seeded to None, but only when nothing has already + # established one. A stream that assembles into a response + # object recomputes the cost right after this; a pass-through + # stream cannot (its body is opaque) and carries the cost its + # upstream reported in the response headers, which an + # unconditional reset would discard. + self.model_call_details.setdefault("response_cost", None) + return None + if not self._is_recognized_call_type_for_logging(logging_result=logging_result) and not isinstance( + logging_result, (dict, list) + ): + return None + return self._process_hidden_params_and_response_cost( + logging_result=logging_result, + start_time=start_time, + end_time=end_time, + build_logging_payload=build_logging_payload, + publish_logging_payload=publish_logging_payload, + ) + def _is_recognized_call_type_for_logging( self, logging_result: object, @@ -2883,11 +2973,14 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: self.success_handler(result=complete_streaming_response) - async def async_flush_passthrough_collected_chunks(self, collector: "PassthroughStreamCollector"): + async def async_flush_passthrough_collected_chunks(self, collector: "PassthroughStreamCollector") -> bool: + """Log the response the collector built from the stream, and say whether there was one to log.""" complete_streaming_response: Final = collector.build_logged_response(litellm_logging_obj=self) - if complete_streaming_response is not None: - await self.async_success_handler(result=complete_streaming_response) + if complete_streaming_response is None: + return False + await self.async_success_handler(result=complete_streaming_response) + return True def _restore_correlation_context(self) -> None: """Restore trace_id/session_id contextvars to their pre-call value. @@ -2931,6 +3024,8 @@ class Logging(LiteLLMLoggingBaseClass): start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, cache_hit: bool | None = None, + *, + for_async_request: bool = False, **kwargs: object, # kwargs-ok: forwarded to _success_handler_body ) -> None: """Restores trace_id/session_id contextvars once this attempt's own success @@ -2938,7 +3033,12 @@ class Logging(LiteLLMLoggingBaseClass): try: with post_response_phase(): return self._success_handler_body( - result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + for_async_request=for_async_request, + **kwargs, ) finally: self._restore_correlation_context() @@ -2949,17 +3049,20 @@ class Logging(LiteLLMLoggingBaseClass): start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, cache_hit: bool | None = None, + for_async_request: bool = False, **kwargs: Any, # kwargs-ok: forwarded from success_handler ) -> None: verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit) if not self.should_run_logging(event_type="sync_success"): # prevent double logging return - start_time, end_time, result = self._success_handler_helper_fn( + start_time, end_time, result, request_logging_payload = self._prepare_success_logging( start_time=start_time, end_time=end_time, result=result, cache_hit=cache_hit, standard_logging_object=kwargs.get("standard_logging_object", None), + build_logging_payload=True, + publish_logging_payload=not for_async_request, ) litellm_params = self.model_call_details.get("litellm_params", {}) is_sync_request: Final = self.is_sync_litellm_request(litellm_params) @@ -2991,16 +3094,45 @@ class Logging(LiteLLMLoggingBaseClass): ) self.model_call_details["complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - logged_streaming_response, start_time, end_time + if not for_async_request: + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + logged_streaming_response, start_time, end_time + ) + standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( + "standard_logging_object" + ) + if standard_logging_payload is not None: + # Only emit for sync requests (async_success_handler handles async) + if is_sync_request: + emit_standard_logging_payload(standard_logging_payload) + callback_logging_payload: Final = ( + self._build_standard_logging_payload( + self.model_call_details["complete_streaming_response"], + start_time, + end_time, + isolate_call_details=True, ) - standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( - "standard_logging_object" + if for_async_request and complete_streaming_response is not None + else request_logging_payload + ) + shared_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( + "standard_logging_object" + ) + isolated_logging_payload: Final = ( + callback_logging_payload + if callback_logging_payload is not None or shared_logging_payload is None + else shared_logging_payload.copy() + ) + call_details: dict[str, Any] # mutable-ok: legacy loggers take a plain dict + call_details = ( # rebind-ok: logging hooks may return a new dict + self.model_call_details + if not for_async_request + else ( + dict(self.model_call_details) + if isolated_logging_payload is None + else {**self.model_call_details, "standard_logging_object": isolated_logging_payload} ) - if standard_logging_payload is not None: - # Only emit for sync requests (async_success_handler handles async) - if is_sync_request: - emit_standard_logging_payload(standard_logging_payload) + ) callbacks: Final = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, @@ -3008,7 +3140,7 @@ class Logging(LiteLLMLoggingBaseClass): ## REDACT MESSAGES ## result = redact_message_input_output_from_logging( - model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), + model_call_details=call_details, result=result, ) ## LOGGING HOOK ## @@ -3018,7 +3150,7 @@ class Logging(LiteLLMLoggingBaseClass): if ( callback.should_run_guardrail( - data=self.model_call_details, + data=call_details, event_type=GuardrailEventHooks.logging_only, ) is not True @@ -3026,18 +3158,21 @@ class Logging(LiteLLMLoggingBaseClass): continue self.shadow_eval_request_snapshot = None - self.model_call_details, result = callback.logging_hook( - kwargs=self.model_call_details, + call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict + kwargs=call_details, result=result, call_type=self.call_type, ) elif isinstance(callback, CustomLogger): - self.model_call_details, result = callback.logging_hook( - kwargs=self.model_call_details, + call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict + kwargs=call_details, result=result, call_type=self.call_type, ) + if not for_async_request: + self.model_call_details = call_details + call_details_view: Final = MappingProxyType[str, object](call_details) self.mark_logging_complete(event_type="sync_success") for callback in callbacks: try: @@ -3051,7 +3186,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "promptlayer" and promptLayerLogger is not None: print_verbose("reaches promptlayer for logging!") promptLayerLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3059,7 +3194,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "supabase" and supabaseClient is not None: print_verbose("reaches supabase for logging!") - kwargs = self.model_call_details + kwargs = call_details # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -3067,10 +3202,10 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches supabase for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] - model = self._model_call_details_view["model"] - messages = self._model_call_details_view["messages"] + model = call_details_view["model"] + messages = call_details_view["messages"] optional_params = kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) supabaseClient.log_event( @@ -3090,7 +3225,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "wandb" and weightsBiasesLogger is not None: print_verbose("reaches wandb for logging!") weightsBiasesLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3099,7 +3234,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "logfire" and logfireLogger is not None: verbose_logger.debug("reaches logfire for success logging!") kwargs = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v @@ -3109,10 +3244,10 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches logfire for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] logfireLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3123,7 +3258,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "lunary" and lunaryLogger is not None: print_verbose("reaches lunary for logging!") model = self.model - kwargs = self.model_call_details + kwargs = call_details input = kwargs.get("messages", kwargs.get("input", None)) @@ -3134,7 +3269,7 @@ class Logging(LiteLLMLoggingBaseClass): if "complete_streaming_response" not in kwargs: continue else: - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] lunaryLogger.log_event( type=type, @@ -3144,7 +3279,7 @@ class Logging(LiteLLMLoggingBaseClass): input=input, user_id=kwargs.get("user", None), # user_props=self.model_call_details.get("user_props", None), - extra=self._model_call_details_view.get("optional_params", {}), + extra=call_details_view.get("optional_params", {}), response_obj=result, start_time=start_time, end_time=end_time, @@ -3154,8 +3289,8 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "helicone" and heliconeLogger is not None: print_verbose("reaches helicone for logging!") model = self.model - messages = self._model_call_details_view["input"] - kwargs = self.model_call_details + messages = call_details_view["input"] + kwargs = call_details # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -3163,7 +3298,7 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches helicone for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] heliconeLogger.log_success( model=model, @@ -3178,7 +3313,7 @@ class Logging(LiteLLMLoggingBaseClass): global langFuseLogger print_verbose("reaches langfuse for success logging!") kwargs = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends @@ -3191,7 +3326,7 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches langfuse for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( globalLangfuseLogger=langFuseLogger, @@ -3216,7 +3351,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "greenscale" and greenscaleLogger is not None: kwargs = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends @@ -3229,7 +3364,7 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches greenscale for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] greenscaleLogger.log_event( kwargs=kwargs, @@ -3240,7 +3375,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "athina" and athinaLogger is not None: deep_copy = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): deep_copy[k] = v athinaLogger.log_event( kwargs=deep_copy, @@ -3251,7 +3386,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "traceloop": deep_copy = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": deep_copy[k] = v traceloopLogger.log_event( @@ -3267,11 +3402,11 @@ class Logging(LiteLLMLoggingBaseClass): if s3Logger is None: s3Logger = S3Logger() if self.stream: - if "complete_streaming_response" in self.model_call_details: + if "complete_streaming_response" in call_details: print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") s3Logger.log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details["complete_streaming_response"], + kwargs=call_details, + response_obj=call_details["complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, @@ -3280,7 +3415,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("S3Logger Logger: Got Stream Event - No complete stream response as yet") else: s3Logger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3294,19 +3429,17 @@ class Logging(LiteLLMLoggingBaseClass): openMeterLogger = OpenMeterLogger() if self.stream and complete_streaming_response is None: openMeterLogger.log_stream_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = self.model_call_details.get( - "complete_streaming_response", {} - ) - result = self._model_call_details_view["complete_response"] + call_details["complete_response"] = call_details.get("complete_streaming_response", {}) + result = call_details_view["complete_response"] openMeterLogger.log_success_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3319,20 +3452,18 @@ class Logging(LiteLLMLoggingBaseClass): ): # custom logger class if self.stream and complete_streaming_response is None: callback.log_stream_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = self.model_call_details.get( - "complete_streaming_response", {} - ) - result = self._model_call_details_view["complete_response"] + call_details["complete_response"] = call_details.get("complete_streaming_response", {}) + result = call_details_view["complete_response"] callback.log_success_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3343,7 +3474,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose(f"success callbacks: Running Custom Callback Function - {callback}") customLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -4184,8 +4315,18 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, cache_hit, + for_async_request=True, ) + def handle_sync_cache_hit_success_callbacks( + self, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + cache_hit: bool | None = None, + ) -> None: + executor.submit(self.success_handler, result, start_time, end_time, cache_hit) + def _should_run_sync_callbacks_for_async_calls(self) -> bool: """ Returns: @@ -4197,7 +4338,10 @@ class Logging(LiteLLMLoggingBaseClass): ) _filtered_success_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks) _filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks) - return len(_filtered_success_callbacks) > 0 + return any( + not (isinstance(callback, str) and callback in INTERNAL_CALLBACK_NAMES) + for callback in _filtered_success_callbacks # pyright: ignore[reportUnknownVariableType] # callback lists are untyped + ) def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool: """ @@ -6805,6 +6949,36 @@ def _autorouter_savings_for_payload( return None +_CACHE_HIT_ZEROED_AMOUNTS: Final = MappingProxyType( + { + "cache_read_cost": CostBreakdown(cache_read_cost=0.0), + "cache_creation_cost": CostBreakdown(cache_creation_cost=0.0), + "reasoning_cost": CostBreakdown(reasoning_cost=0.0), + "original_cost": CostBreakdown(original_cost=0.0), + "discount_amount": CostBreakdown(discount_amount=0.0), + "margin_fixed_amount": CostBreakdown(margin_fixed_amount=0.0), + "margin_total_amount": CostBreakdown(margin_total_amount=0.0), + } +) + + +def _cache_hit_cost_breakdown(served: CostBreakdown | None) -> CostBreakdown: + """Every amount at 0, keeping the pricing basis and rates a fresh call would have had, which savings reports read.""" + zero_costs: Final = CostBreakdown(input_cost=0.0, output_cost=0.0, total_cost=0.0, tool_usage_cost=0.0) + if served is None: + return zero_costs + zeroed_additional_costs: Final = ( + CostBreakdown(additional_costs=dict.fromkeys(served["additional_costs"], 0.0)) + if "additional_costs" in served + else CostBreakdown() + ) + return functools.reduce( + lambda breakdown, zeroed: breakdown | zeroed, + (_CACHE_HIT_ZEROED_AMOUNTS[key] for key in served if key in _CACHE_HIT_ZEROED_AMOUNTS), + served | zero_costs | zeroed_additional_costs, + ) + + def get_standard_logging_object_payload( kwargs: dict | None, init_response_obj: object, @@ -6888,17 +7062,10 @@ def get_standard_logging_object_payload( "user", None ) # maintain backwards compatibility with old request body check - saved_cache_cost: float = 0.0 if cache_hit is True or served_from_cache(init_response_obj): id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id - if cache_hit is True: - saved_cache_cost = ( - logging_obj.response_cost_calculator( - result=init_response_obj, - cache_hit=False, - ) - or 0.0 - ) + uncached_price: Final = logging_obj.uncached_price(result=init_response_obj) if cache_hit is True else None + saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0 ## Get model cost information ## base_model = get_base_model_from_metadata(model_call_details=kwargs) @@ -6915,7 +7082,7 @@ def get_standard_logging_object_payload( # clean up litellm hidden params clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params) - if clean_hidden_params["response_cost"] is None and raw_response_cost is not None: + if cache_hit is True or (clean_hidden_params["response_cost"] is None and raw_response_cost is not None): clean_hidden_params["response_cost"] = llm_response_cost if clean_hidden_params["litellm_overhead_time_ms"] is None and status == "success": # /v1/messages dict results and the bridge stream wrappers keep it on the logging object; @@ -6977,7 +7144,14 @@ def get_standard_logging_object_payload( ): model_name = response_model_name - request_cost_breakdown: Final = cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost) + request_cost_breakdown: Final = cost_breakdown_with_guardrail( + ( + _cache_hit_cost_breakdown(uncached_price.cost_breakdown) + if uncached_price is not None + else logging_obj.cost_breakdown + ), + guardrail_cost, + ) captured_baseline: Final = logging_obj.baseline_observation autorouter_savings: Final = ( None diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index f3c60679c81..8b94c34942c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -251,14 +251,7 @@ async def convert_to_streaming_response_async( setattr( model_response_object, "usage", - Usage( - completion_tokens=response_object["usage"].get("completion_tokens", 0), - prompt_tokens=response_object["usage"].get("prompt_tokens", 0), - total_tokens=response_object["usage"].get("total_tokens", 0), - completion_tokens_details=_cached_completion_tokens_details( - _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) - ), - ), + _replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])), ) if "id" in response_object: @@ -345,12 +338,10 @@ def convert_to_streaming_response( model_response_object.choices = choice_list if "usage" in response_object and response_object["usage"] is not None: - setattr(model_response_object, "usage", Usage()) - model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) - model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) - model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) - model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details( - _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + setattr( + model_response_object, + "usage", + _replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])), ) if "id" in response_object: @@ -414,13 +405,28 @@ class _CachedMessageProviderFields(TypedDict, total=False): class _CachedUsageDetails(TypedDict, total=False): + completion_tokens: ReadOnly[object] + prompt_tokens: ReadOnly[object] + total_tokens: ReadOnly[object] completion_tokens_details: ReadOnly[object] + prompt_tokens_details: ReadOnly[object] + cache_read_input_tokens: ReadOnly[object] + cache_creation_input_tokens: ReadOnly[object] + + +class _PromptCacheTokenCounts(TypedDict, total=False): + cache_read_input_tokens: ReadOnly[int] + cache_creation_input_tokens: ReadOnly[int] _CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields) _CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails) _FIELD_MAP: Final = TypeAdapter(dict[str, object]) _COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper) +_PROMPT_TOKENS_DETAILS: Final = TypeAdapter(PromptTokensDetailsWrapper) +_PROMPT_CACHE_TOKEN_COUNTS: Final = TypeAdapter(_PromptCacheTokenCounts) +_TOKEN_COUNT: Final = TypeAdapter(int) +_PROMPT_CACHE_TOKEN_FIELDS: Final = ("cache_read_input_tokens", "cache_creation_input_tokens") def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None: @@ -438,6 +444,44 @@ def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionT return None +def _cached_prompt_tokens_details(usage: _CachedUsageDetails) -> PromptTokensDetailsWrapper | None: + details: Final = usage.get("prompt_tokens_details") + if not isinstance(details, dict): + return None + try: + return _PROMPT_TOKENS_DETAILS.validate_python(details) + except ValidationError: + return None + + +def _cached_token_count(count: object) -> int | None: + try: + return _TOKEN_COUNT.validate_python(count) + except ValidationError: + return None + + +def _cached_prompt_cache_token_counts(usage: _CachedUsageDetails) -> _PromptCacheTokenCounts: + return _PROMPT_CACHE_TOKEN_COUNTS.validate_python( + { + field: count + for field in _PROMPT_CACHE_TOKEN_FIELDS + if (count := _cached_token_count(usage.get(field))) is not None + } + ) + + +def _replayed_usage(usage: _CachedUsageDetails) -> Usage: + return Usage( + completion_tokens=_cached_token_count(usage.get("completion_tokens")) or 0, + prompt_tokens=_cached_token_count(usage.get("prompt_tokens")) or 0, + total_tokens=_cached_token_count(usage.get("total_tokens")) or 0, + completion_tokens_details=_cached_completion_tokens_details(usage), + prompt_tokens_details=_cached_prompt_tokens_details(usage), + **_cached_prompt_cache_token_counts(usage), + ) + + def handle_invalid_parallel_tool_calls( tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall], ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None: diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 044123d8729..b704dfa2675 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -9,7 +9,8 @@ import json import mimetypes import re from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence -from itertools import groupby, islice +from graphlib import CycleError, TopologicalSorter +from itertools import chain, groupby, islice from os import PathLike from pathlib import Path from types import MappingProxyType @@ -1426,6 +1427,123 @@ def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str, frontier = tuple(child for node in frontier for child in _subschemas(node)) +def schema_fragment_ref(ref: str) -> str: + return f"#{ref}" if ref.startswith(("/$defs/", "/definitions/")) else ref + + +def schema_has_unresolvable_ref(schema: Mapping[str, object]) -> bool: + definitions: Final = MappingProxyType(dict(_root_definitions(schema))) + if not _refs_in(schema) <= definitions.keys(): + return True + try: + tuple(TopologicalSorter({pointer: _refs_in(body) for pointer, body in definitions.items()}).static_order()) + except CycleError: + return True + return False + + +def schema_union_property_count(schema: Mapping[str, object]) -> int: + return sum(1 for _, node in _property_schemas(schema) if _is_union(node)) + + +def schema_optional_property_count(schema: Mapping[str, object]) -> int: + return sum(len(_optional_property_names(node)) for node in _schema_nodes(schema)) + + +def schema_property_count(schema: Mapping[str, object]) -> int: + return sum(len(properties) for node in _schema_nodes(schema) if (properties := _schema_at(node, "properties"))) + + +def schema_grammar_branch_count(schema: Mapping[str, object]) -> int: + return sum(_grammar_branches(node) for node in _schema_nodes(schema)) + + +def schema_has_open_object(schema: Mapping[str, object]) -> bool: + return any(_is_open_object(node) for node in _schema_nodes(schema)) + + +def schema_has_lookaround_pattern(schema: Mapping[str, object]) -> bool: + return any( + isinstance(pattern := node.get("pattern"), str) and _uses_regex_lookaround(pattern) + for node in _schema_nodes(schema) + ) + + +def schema_string_formats(schema: Mapping[str, object]) -> frozenset[str]: + return frozenset( + value for node in _schema_nodes(schema) if _may_be_string(node) and isinstance(value := node.get("format"), str) + ) + + +def _may_be_string(node: Mapping[str, object]) -> bool: + node_type: Final = node.get("type") + return node_type is None or node_type == "string" or "string" in _list_at(node, "type") + + +def is_object_schema(node: Mapping[str, object]) -> bool: + node_type: Final = node.get("type") + return node_type == "object" or (isinstance(node_type, list) and "object" in node_type) + + +def _schema_nodes(schema: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + return chain.from_iterable(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING)) + + +_SCHEMA_LIST: Final = TypeAdapter(tuple[object, ...]) + + +def _list_at(container: Mapping[str, object], key: str) -> tuple[object, ...]: + value: Final = container.get(key) + return _SCHEMA_LIST.validate_python(value) if isinstance(value, list) else () + + +def _named_schemas(container: Mapping[str, object] | None) -> Iterator[tuple[str, Mapping[str, object]]]: + if container is None: + return iter(()) + return ((key, value) for key in container if (value := _schema_at(container, key)) is not None) + + +def _root_definitions(schema: Mapping[str, object]) -> Iterator[tuple[str, Mapping[str, object]]]: + for defs_key in ("$defs", "definitions"): + yield from ((f"#/{defs_key}/{name}", body) for name, body in _named_schemas(_schema_at(schema, defs_key))) + + +def _refs_in(schema: Mapping[str, object]) -> frozenset[str]: + return frozenset( + schema_fragment_ref(ref) for node in _schema_nodes(schema) if isinstance(ref := node.get("$ref"), str) + ) + + +def _property_schemas(schema: Mapping[str, object]) -> Iterator[tuple[str, Mapping[str, object]]]: + return chain.from_iterable(_named_schemas(_schema_at(node, "properties")) for node in _schema_nodes(schema)) + + +def _is_union(node: Mapping[str, object]) -> bool: + return isinstance(node.get("anyOf"), list) or isinstance(node.get("oneOf"), list) or len(_list_at(node, "type")) > 1 + + +def _grammar_branches(node: Mapping[str, object]) -> int: + properties: Final = _schema_at(node, "properties") + if properties is None: + return 0 + optional: Final = len(_optional_property_names(node)) + return 2**optional * (len(properties) - optional + 1) + + +def _optional_property_names(node: Mapping[str, object]) -> frozenset[str]: + properties: Final = _schema_at(node, "properties") + if properties is None: + return frozenset() + return frozenset(properties) - frozenset(name for name in _list_at(node, "required") if isinstance(name, str)) + + +def _is_open_object(node: Mapping[str, object]) -> bool: + additional_properties: Final = node.get("additionalProperties") + if additional_properties is True or isinstance(additional_properties, dict) or "patternProperties" in node: + return True + return is_object_schema(node) and "properties" not in node and additional_properties is not False + + def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]: for key, value in node.items(): if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 87e76fe39f7..0e9206b422e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5504,14 +5504,14 @@ def _bedrock_tools_pt( description = name defs = parameters.pop("$defs", {}) - defs_copy = copy.deepcopy(defs) + defs_copy = copy.deepcopy(defs) if isinstance(defs, dict) else {} # Expand $ref references in parameters using the definitions # Note: We don't pre-flatten defs as that causes exponential memory growth # with circular references (see issue #19098). unpack_defs handles nested # refs recursively and correctly detects/skips circular references. unpack_defs(parameters, defs_copy) normalize_json_schema_custom_types_to_object(parameters) - if parameters.get("type") not in _valid_json_schema_root_types: + if not (isinstance(root_type := parameters.get("type"), str) and root_type in _valid_json_schema_root_types): parameters["type"] = "object" tool_block = cast( BedrockToolBlock, diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index e4eb9acc37d..927caec3350 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,11 +1,11 @@ import asyncio import json import traceback -from collections.abc import Callable, Coroutine, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, MutableMapping, Sequence from dataclasses import dataclass from enum import Enum, auto from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast, runtime_checkable from pydantic import TypeAdapter from typing_extensions import ReadOnly @@ -14,6 +14,7 @@ import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger from litellm._uuid import uuid from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -100,6 +101,16 @@ class ScopedWebSocket(Protocol): def scope(self) -> _ASGIScope: ... +@runtime_checkable +class _ReplayScopeWebSocket(Protocol): + @property + def scope(self) -> MutableMapping[str, object]: ... + + +def client_websocket_scope(websocket: object) -> MutableMapping[str, object] | None: + return websocket.scope if isinstance(websocket, _ReplayScopeWebSocket) else None + + class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... @@ -233,6 +244,7 @@ class RealTimeStreaming: logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket + self._replay_scope: Final = client_websocket_scope(websocket) self.backend_ws = backend_ws self.logging_obj = logging_obj self._logging_worker = logging_worker @@ -1451,6 +1463,13 @@ class RealTimeStreaming: return item async def _receive_client_message(self) -> str | None: + refused_opening_message: Final = ( + None + if self._replay_scope is None + else self._replay_scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + ) + if isinstance(refused_opening_message, str): + return refused_opening_message try: return await self.websocket.receive_text() except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 78fb82c98e7..78cfb58bedf 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -47,7 +47,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..exceptions import OpenAIError from .core_helpers import map_finish_reason, process_response_headers -from .exception_mapping_utils import exception_type +from .exception_mapping_utils import exception_type, with_deployment_api_base from .llm_response_utils.get_api_base import get_api_base from .rules import Rules @@ -1579,11 +1579,7 @@ class CustomStreamWrapper: except Exception as e: traceback.format_exc() setattr(e, "message", str(e)) - raise exception_type( - model=self.model, - custom_llm_provider=self.custom_llm_provider, - original_exception=e, - ) + raise self._map_stream_exception(e) def set_logging_event_loop(self, loop): """ @@ -2370,6 +2366,14 @@ class CustomStreamWrapper: if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): self._emitted_disqualifying_content = True + def _map_stream_exception(self, original_exception: Exception) -> Exception: + return exception_type( + model=self.model, + custom_llm_provider=self.custom_llm_provider, + original_exception=original_exception, + extra_kwargs=with_deployment_api_base(self.logging_obj.model_call_details.get("litellm_params")), + ) + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. @@ -2414,13 +2418,7 @@ class CustomStreamWrapper: mapped_exception: Exception = e else: try: - mapped_exception = exception_type( - model=self.model, - custom_llm_provider=self.custom_llm_provider, - original_exception=e, - completion_kwargs={}, - extra_kwargs={}, - ) + mapped_exception = self._map_stream_exception(e) except Exception as mapping_error: mapped_exception = mapping_error self._restore_consumer_correlation_context() diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6b6974becac..7654b56ad9a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -30,6 +30,7 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) +from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details from litellm.llms.anthropic.pass_through.stream_assembly import ( build_complete_streaming_response, build_usage_only_response_from_chunks, @@ -1456,6 +1457,9 @@ class AnthropicMessagesHandler(BaseTranslation): return _SSEFieldRewrite("delta", "text", next(replacements)) cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + cls._redact_ended_stream_message_delta( + responses_so_far, kept_tool_use_ids=frozenset(cls._streamed_tool_use_ids(responses_so_far).values()) + ) @classmethod def _write_ended_stream_tool_call_rewrites( @@ -1523,6 +1527,64 @@ class AnthropicMessagesHandler(BaseTranslation): return None cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) + streamed_tool_use_ids: Final = cls._streamed_tool_use_ids(responses_so_far) + cls._redact_ended_stream_message_delta( + responses_so_far, + kept_tool_use_ids=frozenset( + tool_use_id for index, tool_use_id in streamed_tool_use_ids.items() if index not in rewrites_by_block + ), + ) + + @classmethod + def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]: + """The id of each buffered ``tool_use`` block, by its content block index.""" + events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far) + tool_use_blocks: Final = ( + (event.get("index"), block) + for event in events + if event.get("type") == "content_block_start" + and (block := as_str_mapping(event.get("content_block"))) is not None + and block.get("type") == "tool_use" + ) + return MappingProxyType( + { + index: tool_use_id + for index, block in tool_use_blocks + if isinstance(index, int) and isinstance(tool_use_id := block.get("id"), str) + } + ) + + @classmethod + def _redact_ended_stream_message_delta( + cls, + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + *, + kept_tool_use_ids: frozenset[str], + ) -> None: + """Hold the buffered ``message_delta`` to the rules a re-emitted rewrite follows: a verdict survives + only for a tool call delivered unchanged, and no free-text ``explanation`` describes the text the + guardrail rewrote.""" + + def rewrite_stop_details(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None + if delta is None or "stop_details" not in delta: + return None + return _SSEFieldRewrite( + "delta", "stop_details", redacted_stop_details(delta["stop_details"], keep_explanations=False) + ) + + def rewrite_safeguard_results(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None + if delta is None or "safeguard_results" not in delta: + return None + return _SSEFieldRewrite( + "delta", + "safeguard_results", + filtered_safeguard_results(delta["safeguard_results"], kept_tool_use_ids, keep_explanations=False), + ) + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_stop_details) + cls._rewrite_ended_stream_events(responses_so_far, rewrite_safeguard_results) @staticmethod def _rewrite_ended_stream_events( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0b2f75e352b..38af6e7998b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1,3 +1,4 @@ +import copy import json import re import time @@ -24,7 +25,18 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( + DEFS_MAX_INLINED_BYTES, + is_object_schema, sanitize_input_schema_for_anthropic, + schema_fragment_ref, + schema_grammar_branch_count, + schema_has_lookaround_pattern, + schema_has_open_object, + schema_has_unresolvable_ref, + schema_optional_property_count, + schema_string_formats, + schema_union_property_count, + unpack_defs, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( RemoteMedia, @@ -43,7 +55,7 @@ from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, unwrap_response_format_envelope, ) -from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS, BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_BETA_HEADER_VALUES, @@ -150,6 +162,12 @@ _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 _DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool) _TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _TOOLS_ADAPTER: Final = TypeAdapter(list[object]) +_NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset( + ("date-time", "time", "date", "duration", "email", "hostname", "uri", "ipv4", "ipv6", "uuid") +) +NATIVE_OUTPUT_MAX_UNION_PARAMETERS: Final = 16 +NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS: Final = 24 +NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES: Final = 2048 def _tool_type(tool: object) -> str | None: @@ -753,7 +771,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): grammar Anthropic compiles the schema into, so the API rejects them with a 400 ``invalid_request_error`` (e.g. "output_format.schema: For 'array' type, property 'uniqueItems' is not supported"): - - maxItems/minItems/uniqueItems/contains/minContains/maxContains/prefixItems: array constraints + - maxItems/minItems above 1/uniqueItems/contains/minContains/maxContains/prefixItems: array constraints - minimum/maximum/exclusiveMinimum/exclusiveMaximum/multipleOf: numeric constraints - minLength/maxLength: string constraints - minProperties/maxProperties/patternProperties/propertyNames: object constraints @@ -810,7 +828,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "else": "conditional (else): {}", "not": "must not match: {}", } - unsupported_fields: Final = set(constraint_labels) + min_items: Final = schema.get("minItems") + enforced_min_items: Final = ( + isinstance(min_items, int) and not isinstance(min_items, bool) and min_items in (0, 1) + ) + unsupported_fields: Final = set(constraint_labels) - ({"minItems"} if enforced_min_items else set()) # Build description additions from removed constraints. Iterating # constraint_labels (not the set) keeps the note order deterministic across @@ -818,7 +840,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # PYTHONHASHSEED and stay cache-friendly. constraint_descriptions: Final[list] = [] for field, label in constraint_labels.items(): - if field not in schema: + if field not in schema or field not in unsupported_fields: continue value = schema[field] # A falsy boolean constraint (e.g. ``uniqueItems: false``) imposes no @@ -853,7 +875,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Already handled above continue - if key in ("properties", "$defs") and isinstance(value, dict): + if key in ("properties", "$defs", "definitions") and isinstance(value, dict): nested_schemas: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in nested_schemas.items()} elif key == "items" and isinstance(value, dict): @@ -865,6 +887,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ] elif key == "allOf" and isinstance(value, list): result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] + elif key == "$ref" and isinstance(value, str): + result[key] = schema_fragment_ref(value) elif key == "oneOf" and isinstance(value, list): # Anthropic rejects oneOf ("Schema type 'oneOf' is not supported"); # the Anthropic SDK rewrites it to anyOf, so do the same. @@ -876,7 +900,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Anthropic requires additionalProperties=false for object schemas # See: https://docs.anthropic.com/en/docs/build-with-claude/structured-outputs - if result.get("type") == "object": + if is_object_schema(result): result["additionalProperties"] = False return result @@ -1628,32 +1652,55 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return json_schema + @staticmethod + def _defs_unpacked(json_schema: Mapping[str, object]) -> Mapping[str, object]: + unpacked: Final = copy.deepcopy(dict(json_schema)) + defs: Final = unpacked.pop("$defs", unpacked.pop("definitions", None)) + if isinstance(defs, dict) and defs: + unpack_defs(unpacked, defs, max_inlined_bytes=DEFS_MAX_INLINED_BYTES) + return unpacked + + @staticmethod + def native_output_schema(json_schema: Mapping[str, object]) -> Mapping[str, object] | None: + if ( + schema_has_unresolvable_ref(json_schema) + or schema_has_open_object(json_schema) + or schema_has_lookaround_pattern(json_schema) + or not schema_string_formats(json_schema) <= _NATIVE_OUTPUT_STRING_FORMATS + or schema_union_property_count(json_schema) > NATIVE_OUTPUT_MAX_UNION_PARAMETERS + or schema_optional_property_count(json_schema) > NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS + or schema_grammar_branch_count(json_schema) > NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES + ): + return None + return AnthropicConfig.filter_anthropic_output_schema(dict(json_schema)) + + def uses_native_structured_output_for_schema( + self, + model: str, + json_schema: Mapping[str, object] | None, + request_params: Mapping[str, object], + ) -> bool: + open_object_takes_forced_json_tool: Final = ( + json_schema is not None + and schema_has_open_object(json_schema) + and not self.is_thinking_requested(request_params) + and not self.is_thinking_enabled(dict(request_params)) + and not request_params.get("tools") + and not AnthropicModelInfo.forced_tool_use_unsupported(model) + ) + return self.uses_native_structured_output(model) and not open_object_takes_forced_json_tool + + def native_output_format_schema(self, json_schema: Mapping[str, object]) -> Mapping[str, object] | None: + return self.filter_anthropic_output_schema(dict(self._defs_unpacked(json_schema))) + def map_response_format_to_anthropic_output_format(self, value: dict | None) -> AnthropicOutputSchema | None: - json_schema: dict | None = self._extract_json_schema_from_response_format(value) + json_schema: Final = self._extract_json_schema_from_response_format(value) if json_schema is None: return None - - # Resolve $ref/$defs before filtering — Anthropic doesn't support - # external schema references (e.g., /$defs/CalendarEvent). - import copy - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - DEFS_MAX_INLINED_BYTES, - unpack_defs, - ) - - json_schema = copy.deepcopy(json_schema) - defs: Final = json_schema.pop("$defs", json_schema.pop("definitions", {})) - if defs: - unpack_defs(json_schema, defs, max_inlined_bytes=DEFS_MAX_INLINED_BYTES) - - # Filter out unsupported fields for Anthropic's output_format API - filtered_schema: Final = self.filter_anthropic_output_schema(json_schema) - - return AnthropicOutputSchema( - type="json_schema", - schema=filtered_schema, - ) + native_schema: Final = self.native_output_format_schema(json_schema) + if native_schema is None: + return None + return AnthropicOutputSchema(type="json_schema", schema=dict(native_schema)) def map_response_format_to_anthropic_tool( self, value: dict | None, optional_params: dict, is_thinking_enabled: bool @@ -1820,7 +1867,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): output_key=param, ) elif param == "response_format" and isinstance(value, dict): - if self.uses_native_structured_output(model): + if self.uses_native_structured_output_for_schema( + model, + self._extract_json_schema_from_response_format(value), + REQUEST_PARAMS.validate_python(non_default_params), + ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: optional_params["output_format"] = _output_format @@ -2979,6 +3030,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): completion_response["stop_reason"] = "stop" if refusal_fields: json_mode_message.provider_specific_fields = dict(refusal_fields) + if thinking_blocks: + json_mode_message.thinking_blocks = thinking_blocks + json_mode_message.reasoning_content = reasoning_content _message = json_mode_message model_response.choices[0].message = _message diff --git a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py index 5930739d445..8d1b358c8f1 100644 --- a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py @@ -12,6 +12,7 @@ import json from collections.abc import Mapping from typing import Any, Final, cast +from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -34,8 +35,13 @@ class FakeAnthropicMessagesStreamIterator: - message_stop """ - def __init__(self, response: AnthropicMessagesResponse): + def __init__( + self, + response: AnthropicMessagesResponse, + message_delta_fields: Mapping[str, object] | None = None, + ): self.response = response + self.message_delta_fields = message_delta_fields self.chunks = self._create_streaming_chunks() self.current_index = 0 @@ -175,14 +181,9 @@ class FakeAnthropicMessagesStreamIterator: cache_read_input_tokens: Final = usage.get("cache_read_input_tokens") if cache_read_input_tokens is not None: delta_usage["cache_read_input_tokens"] = cache_read_input_tokens - safeguard_results: Final = self.response.get("safeguard_results") message_delta: Final = { "type": "message_delta", - "delta": { - "stop_reason": self.response.get("stop_reason"), - "stop_sequence": self.response.get("stop_sequence"), - **({"safeguard_results": safeguard_results} if safeguard_results is not None else {}), - }, + "delta": anthropic_message_delta(response_dict, self.message_delta_fields), "usage": delta_usage, } chunks.append(f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode()) diff --git a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py index e65fe596da7..44db193ea83 100644 --- a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py @@ -18,7 +18,11 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP -from litellm.llms.anthropic.pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE +from litellm.llms.anthropic.pass_through.messages.utils import ( + INCOMPLETE_STREAM_ERROR_MESSAGE, + MESSAGE_DELTA_OPTIONAL_KEYS, + anthropic_message_delta, +) from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -280,30 +284,32 @@ def anthropic_messages_response_as_sse_events(response: AnthropicMessagesRespons content_events: Final = ( event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block) ) - # A real message_start always carries a null stop_reason/stop_sequence and - # a zero output_tokens - those are only known once generation finishes, so - # copying the completed response's final values here would let a client - # treat the message as already finished, or double-count output tokens. + # A real message_start always carries a null stop_reason/stop_sequence/stop_details, + # a zero output_tokens and no safeguard_results - those are only known once + # generation finishes, so copying the completed response's final values here + # would let a client treat the message as already finished, or double-count + # output tokens. message_start_usage: Final = { **(response.get("usage") or {}), "output_tokens": 0, } + message_start_fields: Final = { + key: value for key, value in response.items() if key not in MESSAGE_DELTA_OPTIONAL_KEYS + } message_start_payload: Final = { "type": "message_start", "message": { - **response, + **message_start_fields, "content": [], "stop_reason": None, "stop_sequence": None, + "stop_details": None, "usage": message_start_usage, }, } message_delta_payload: Final = { "type": "message_delta", - "delta": { - "stop_reason": response.get("stop_reason"), - "stop_sequence": response.get("stop_sequence"), - }, + "delta": anthropic_message_delta(response), "usage": response.get("usage") or {}, } return ( diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index 640736d6463..ad45485fa56 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -4,6 +4,7 @@ from typing import ClassVar, Final import httpx from litellm._logging import verbose_logger +from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -708,6 +709,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params: dict, custom_llm_provider: str = "anthropic", messages: Sequence[object] = (), + beta_map_provider: str | None = None, ) -> dict: """ Auto-inject anthropic-beta headers based on features used. @@ -724,8 +726,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params: Optional parameters including tools, context_management, output_format, speed custom_llm_provider: Provider name for looking up correct tool search header messages: Request messages, scanned for per-message output_config + beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider """ - beta_values: Final[set] = set() + beta_values: Final[set[str]] = set() existing_beta: Final = tuple( piece.strip() @@ -775,6 +778,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if optional_params.get("speed") == "fast": beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value) + if optional_params.get("safeguards") is not None and is_beta_header_supported_by_bundled_config( + ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value, + beta_map_provider or custom_llm_provider, + ): + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + if messages_carry_output_config(messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index d8b96c1b7a0..d91a5e2fb5c 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -2,7 +2,7 @@ from collections.abc import Iterable, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Final, cast, get_type_hints -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import as_str_mapping @@ -49,6 +49,95 @@ def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails: return AnthropicStopDetails(type="refusal", category=None, explanation=explanation) +MESSAGE_DELTA_OPTIONAL_KEYS: Final = ("stop_details", "safeguard_results") + + +def anthropic_message_delta( + response: Mapping[str, object], optional_fields: Mapping[str, object] | None = None +) -> Mapping[str, object]: + """ + The ``delta`` of the ``message_delta`` event synthesized from a complete + Anthropic Messages response. A real stream carries ``stop_details`` and + ``safeguard_results`` there, so a client that asked to stream reads them + from this event and nowhere else. ``optional_fields`` replaces the + response's own when it was rebuilt from a stream that carried them. + """ + optional_source: Final = response if optional_fields is None else optional_fields + present_optional: Final = { + key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := optional_source.get(key)) is not None + } + return { + "stop_reason": response.get("stop_reason"), + "stop_sequence": response.get("stop_sequence"), + **present_optional, + } + + +_SAFEGUARD_RESULT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def redacted_stop_details(raw: object, keep_explanations: bool) -> object: + """``stop_details`` as a rewritten stream may deliver it: its free-text ``explanation`` can quote + the text a guardrail rewrote, so it is nulled unless the rewrite changed no text.""" + stop_details: Final = as_str_mapping(raw) + if keep_explanations or stop_details is None or "explanation" not in stop_details: + return raw + return {**stop_details, "explanation": None} + + +def _without_explanation(raw: object) -> object: + mapping: Final = as_str_mapping(raw) + if mapping is None: + return raw + return {key: value for key, value in mapping.items() if key != "explanation"} + + +def _status_without_explanations(raw: object) -> object: + status: Final = as_str_mapping(_without_explanation(raw)) + if status is None: + return raw + tool_uses: Final = as_str_mapping(status.get("tool_uses")) + if tool_uses is None: + return status + return { + **status, + "tool_uses": {tool_use_id: _without_explanation(verdict) for tool_use_id, verdict in tool_uses.items()}, + } + + +def _result_without_explanations(raw: object) -> object: + result: Final = as_str_mapping(_without_explanation(raw)) + if result is None: + return raw + if "status" not in result: + return result + return {**result, "status": _status_without_explanations(result["status"])} + + +def _kept_tool_use_verdicts(raw: object, kept_tool_use_ids: frozenset[str]) -> object: + result: Final = as_str_mapping(raw) + status: Final = as_str_mapping(result.get("status")) if result is not None else None + tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None + if result is None or status is None or tool_uses is None: + return raw + kept_tool_uses: Final = { + tool_use_id: verdict for tool_use_id, verdict in tool_uses.items() if tool_use_id in kept_tool_use_ids + } + return {**result, "status": {**status, "tool_uses": kept_tool_uses}} + + +def filtered_safeguard_results(raw: object, kept_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: + """``safeguard_results`` as a rewritten stream may deliver them: a ``tool_uses`` verdict survives only + for a tool call the client gets with the input the classifier judged, and the free-text ``explanation`` of a + result, its status and each verdict is dropped unless the rewrite changed no text.""" + try: + results: Final = _SAFEGUARD_RESULT_ITEMS.validate_python(raw) + except ValidationError: + return raw if keep_explanations else None + kept: Final = tuple(_kept_tool_use_verdicts(result, kept_tool_use_ids) for result in results) + return kept if keep_explanations else tuple(_result_without_explanations(result) for result in kept) + + def _mapping_field(container: object, key: str) -> object | None: """One key of a raw provider payload, or None when the payload is not a mapping.""" if not isinstance(container, Mapping): diff --git a/litellm/llms/anthropic/pass_through/stream_assembly.py b/litellm/llms/anthropic/pass_through/stream_assembly.py index 89961bbe255..90cb9f4c335 100644 --- a/litellm/llms/anthropic/pass_through/stream_assembly.py +++ b/litellm/llms/anthropic/pass_through/stream_assembly.py @@ -1,16 +1,18 @@ import codecs import json from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import as_str_mapping, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.pass_through.messages.utils import MESSAGE_DELTA_OPTIONAL_KEYS from litellm.types.utils import ( Choices, Message, @@ -447,6 +449,27 @@ def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: ) +def anthropic_message_delta_fields(all_chunks: Sequence[object]) -> Mapping[str, object]: + """The ``stop_details`` and ``safeguard_results`` of the stream's ``message_delta``. + + The assembled OpenAI shape has no slot for them, and a client that asked to stream reads them + from that event and nowhere else (a Claude Code auto-mode session takes its classifier verdict + from ``safeguard_results``), so a caller re-emitting a rewritten response merges them back in. + """ + sse_stream: Final = _joined_sse_stream(all_chunks) + if sse_stream is None: + return MappingProxyType({}) + message_delta: Final = next( + (event for event in reversed(_parsed_sse_events(sse_stream)) if event.get("type") == "message_delta"), None + ) + delta: Final = as_str_mapping(message_delta.get("delta")) if message_delta is not None else None + if delta is None: + return MappingProxyType({}) + return MappingProxyType( + {key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := delta.get(key)) is not None} + ) + + def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: """Whether raw SSE frames are Anthropic Messages events. diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index fc22e9a8b73..7aae2cb949d 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -46,14 +46,10 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) -def _is_grok_model(model: str) -> bool: - return "grok" in model - - class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract usage: Final = chunk.get("usage") - if _is_grok_model(str(chunk.get("model") or "")) and isinstance(usage, dict): + if isinstance(usage, dict): XAIChatConfig.fold_reasoning_tokens_into_completion(usage) return super().chunk_parser(chunk) @@ -127,7 +123,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Check if the model supports stop tokens. """ - if _is_grok_model(model): + if "grok" in model: # Reuse Xai method for Grok model xai_config: Final = XAIChatConfig() return xai_config.supports_stop_reason(model) @@ -345,8 +341,7 @@ class AzureAIStudioConfig(OpenAIConfig): api_key=api_key, json_mode=json_mode, ) - if _is_grok_model(model): - XAIChatConfig.fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) return response def get_model_response_iterator( diff --git a/litellm/llms/azure_ai/decisions/transformation.py b/litellm/llms/azure_ai/decisions/transformation.py index 9b10f393c7c..fb5330ed715 100644 --- a/litellm/llms/azure_ai/decisions/transformation.py +++ b/litellm/llms/azure_ai/decisions/transformation.py @@ -1,11 +1,21 @@ +import asyncio import re +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import httpx +from pydantic import ValidationError +import litellm +from litellm.llms.azure_ai.common_utils import ( + AZURE_ENTRA_LITELLM_PARAM_KEYS, + get_azure_ai_auth_headers, +) from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig _FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$") +_AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"}) class AzureAIDecisionsConfig(BaseDecisionsConfig): @@ -14,10 +24,69 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig): path = "/providers/microsoft/v1/systemone" api_key_env = ("AZURE_AI_API_KEY",) api_base_env = ("AZURE_AI_API_BASE",) + api_key_required = False def missing_api_base_message(self, custom_llm_provider: str) -> str: return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://.services.ai.azure.com" + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: + credential_params: Final = MappingProxyType( + {key: litellm_params[key] for key in _AZURE_AI_CREDENTIAL_PARAM_KEYS if key in litellm_params} + ) + try: + auth_headers: Final = get_azure_ai_auth_headers(api_key, credential_params) + except ValidationError as error: + field_names: Final = tuple( + dict.fromkeys( + ".".join(str(field) for field in detail["loc"]) + for detail in error.errors(include_input=False, include_context=False, include_url=False) + ) + ) + raise litellm.BadRequestError( + message=f"Invalid Azure AI credential parameter(s): {', '.join(f'`{name}`' for name in field_names)}", + model=model, + llm_provider="azure_ai", + ) from error + except ValueError as error: + raise litellm.AuthenticationError(message=str(error), model=model, llm_provider="azure_ai") from error + return { + **super().validate_environment( + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ), + **auth_headers, + } + + async def avalidate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: + if api_key: + return self.validate_environment( + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ) + return await asyncio.to_thread( + self.validate_environment, + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ) + def get_complete_url(self, api_base: str, model: str) -> str: url: Final = httpx.URL(api_base) resource_path: Final = _FOUNDRY_ROUTE_SUFFIX.sub("", url.path.rstrip("/")) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index db4d5e0f820..51c8e4b8bef 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Union import httpx -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm.constants import DEFAULT_MAX_TOKENS, RESPONSE_FORMAT_TOOL_NAME from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -38,6 +38,9 @@ else: LiteLLMLoggingObj = Any +REQUEST_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + class BaseLLMException(Exception): def __init__( self, @@ -126,6 +129,17 @@ class BaseConfig(ABC): thinking_type: Final = thinking.get("type") if isinstance(thinking, dict) else None return thinking is True or thinking_type == "enabled" or non_default_params.get("reasoning_effort") is not None + def is_thinking_requested(self, non_default_params: Mapping[str, object]) -> bool: + match non_default_params: + case {"thinking": True | {"type": "enabled" | "adaptive"}}: + return True + case {"reasoning_effort": None | "none"}: + return False + case {"reasoning_effort": object()}: + return True + case _: + return False + def is_max_tokens_in_request(self, non_default_params: dict) -> bool: """ OpenAI spec allows max_tokens or max_completion_tokens to be specified. diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index df7b7645840..3af023df378 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -390,8 +390,24 @@ def systemone_response_to_ir(response: DecisionsResponse, request: DecisionsIRRe ) +_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_LABELS: Final[TypeAdapter[tuple[str, ...]]] = TypeAdapter(tuple[str, ...]) + + +def _normalize_answer(value: object) -> Mapping[str, object]: + answer: Final = _MAPPING.validate_python(value) + legend: Final = answer.get("legend") + if answer.get("type") != "score" or not isinstance(legend, list): + return answer + labels: Final = _LABELS.validate_python(legend) + return {**answer, "legend": {str(index): label for index, label in enumerate(labels)}} + + def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: - return systemone_response_to_ir(_SYSTEMONE_RESPONSE_ADAPTER.validate_python(payload), request) + response: Final = _MAPPING.validate_python(payload) + answers: Final = _MAPPING.validate_python(response.get("answers")) + normalized: Final = {**response, "answers": {name: _normalize_answer(answer) for name, answer in answers.items()}} + return systemone_response_to_ir(_SYSTEMONE_RESPONSE_ADAPTER.validate_python(normalized), request) def _systemone_answer(answer: DecisionsIRAnswer) -> DecisionAnswer | None: @@ -510,13 +526,33 @@ class BaseDecisionsConfig(ABC): def request_model(self, model: str) -> str: return model - def validate_environment(self, headers: Mapping[str, str], model: str, api_key: str | None) -> dict[str, str]: + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: return { **{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS}, **({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}), "Content-Type": "application/json", } + async def avalidate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: + return self.validate_environment( + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ) + def get_complete_url(self, api_base: str, model: str) -> str: return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}" diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 88c4dbfdce6..ee2f6eb9f1a 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -203,6 +203,11 @@ class BaseTranslation(ABC): def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: return None + def needs_ended_stream_scan(self, guardrail: "CustomGuardrail", responses_so_far: Sequence[object]) -> bool: + """Whether the end-of-stream round scans ``guardrail`` even when a sampled round already covered every + text, because this translation hands that guardrail marks only the ended round carries.""" + return False + def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]: """The chunks a client left the stream with, closed the way this endpoint ends a stream, so the end-of-stream scan also inspects tool calls the stream never finished""" diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1f09b3724f1..215db3a4ee4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_lookaround_regex_patterns, + is_object_schema, native_redacted_thinking_data, parse_content_for_reasoning, tool_with_sanitized_parameters, @@ -53,12 +54,21 @@ from litellm.llms.anthropic.chat.transformation import ( tools_without_unsupported_tool_search, ) from litellm.llms.anthropic.common_utils import AnthropicModelInfo -from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS, BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.tool_result_images import place_tool_result_images -from litellm.llms.bedrock.common_utils import bedrock_model_supports_regex_lookaround, bedrock_supports_tool_search +from litellm.llms.bedrock.common_utils import ( + bedrock_chat_uses_native_structured_output, + bedrock_cleaned_output_schema, + bedrock_model_supports_regex_lookaround, + bedrock_native_output_format, + bedrock_supports_tool_search, + deployment_model_info, +) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, - bedrock_request_metadata_is_owned, + bedrock_request_metadata_is_rewritten, + bedrock_request_metadata_violation, + caller_request_metadata_violations_are_dropped, merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) @@ -346,77 +356,15 @@ class AmazonConverseConfig(BaseConfig): and v is not None } - def _validate_request_metadata(self, metadata: dict) -> None: - """ - Validate requestMetadata according to AWS Bedrock Converse API constraints. - - Constraints: - - Maximum of 16 items - - Keys: 1-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{1,256} - - Values: 0-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{0,256} - """ - import re - - if not isinstance(metadata, dict): + def _validate_request_metadata(self, metadata: object) -> None: + violation: Final = bedrock_request_metadata_violation(metadata) + if violation is not None: raise litellm.exceptions.BadRequestError( - message="requestMetadata must be a dictionary", + message=violation.message, model="bedrock", llm_provider="bedrock", ) - if len(metadata) > 16: - raise litellm.exceptions.BadRequestError( - message="requestMetadata can contain a maximum of 16 items", - model="bedrock", - llm_provider="bedrock", - ) - - key_pattern: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$") - value_pattern: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$") - - for key, value in metadata.items(): - if not isinstance(key, str): - raise litellm.exceptions.BadRequestError( - message="requestMetadata keys must be strings", - model="bedrock", - llm_provider="bedrock", - ) - - if not isinstance(value, str): - raise litellm.exceptions.BadRequestError( - message="requestMetadata values must be strings", - model="bedrock", - llm_provider="bedrock", - ) - - if len(key) == 0 or len(key) > 256: - raise litellm.exceptions.BadRequestError( - message="requestMetadata key length must be 1-256 characters", - model="bedrock", - llm_provider="bedrock", - ) - - if len(value) > 256: - raise litellm.exceptions.BadRequestError( - message="requestMetadata value length must be 0-256 characters", - model="bedrock", - llm_provider="bedrock", - ) - - if not key_pattern.match(key): - raise litellm.exceptions.BadRequestError( - message=f"requestMetadata key '{key}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]", - model="bedrock", - llm_provider="bedrock", - ) - - if not value_pattern.match(value): - raise litellm.exceptions.BadRequestError( - message=f"requestMetadata value '{value}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]", - model="bedrock", - llm_provider="bedrock", - ) - @staticmethod def _is_openai_gpt_reasoning_model(model: str) -> bool: return re.search(r"openai\.gpt-\d", model) is not None @@ -939,17 +887,22 @@ class AmazonConverseConfig(BaseConfig): return _tool @staticmethod - def _supports_native_structured_outputs(model: str, custom_llm_provider: str | None = None) -> bool: - """Check if the Bedrock model supports native structured outputs (outputConfig.textFormat). + def _supports_native_structured_outputs( + model: str, + custom_llm_provider: str | None, + thinking_requested: bool, + model_info: Mapping[str, object] | None, + ) -> bool: + """Check if the Bedrock model takes native structured outputs (outputConfig.textFormat) for this request. - Delegates to the standard ``supports_native_structured_output`` utility - which looks up the flag in ``litellm.model_cost`` via - ``_get_model_info_helper``. + ``supports_native_structured_output`` in ``litellm.model_cost`` turns it on for every request, and + ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) only when the + request asks for thinking, which the json tool path would drop. Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html """ - from litellm.utils import supports_native_structured_output - - return supports_native_structured_output(model=model, custom_llm_provider=custom_llm_provider) + return bedrock_chat_uses_native_structured_output( + model, thinking_requested, custom_llm_provider or "bedrock_converse", model_info + ) @staticmethod def _add_additional_properties_to_schema(schema: dict) -> dict: @@ -964,9 +917,9 @@ class AmazonConverseConfig(BaseConfig): if not isinstance(schema, dict): return schema - result: Final = dict(schema) + result: Final[dict[str, object]] = dict(schema) - if result.get("type") == "object" and "additionalProperties" not in result: + if is_object_schema(result) and "additionalProperties" not in result: result["additionalProperties"] = False # Recurse into nested schemas @@ -977,16 +930,46 @@ class AmazonConverseConfig(BaseConfig): if "items" in result and isinstance(result["items"], dict): result["items"] = AmazonConverseConfig._add_additional_properties_to_schema(result["items"]) for defs_key in ("$defs", "definitions"): - if defs_key in result and isinstance(result[defs_key], dict): + if isinstance(definitions := result.get(defs_key), dict): result[defs_key] = { - k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in result[defs_key].items() + k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in definitions.items() } for key in ("anyOf", "allOf", "oneOf"): - if key in result and isinstance(result[key], list): - result[key] = [AmazonConverseConfig._add_additional_properties_to_schema(item) for item in result[key]] + if isinstance(variants := result.get(key), list): + result[key] = [AmazonConverseConfig._add_additional_properties_to_schema(item) for item in variants] return result + def _native_output_schema( + self, + model: str, + json_schema: Mapping[str, object], + thinking_requested: bool, + model_info: Mapping[str, object] | None, + ) -> Mapping[str, object] | None: + if not self._supports_native_structured_outputs( + model, self.custom_llm_provider, thinking_requested, model_info + ): + return None + if BedrockModelInfo.get_base_model(model).startswith("anthropic"): + return bedrock_cleaned_output_schema(model, json_schema, self.custom_llm_provider or "bedrock_converse") + return json_schema + + def _output_config_for_anthropic_format( + self, model: str, output_config_format: Mapping[str, object], model_info: Mapping[str, object] | None + ) -> OutputConfigBlock | None: + native_format: Final = bedrock_native_output_format( + model, output_config_format, self.custom_llm_provider or "bedrock_converse", model_info + ) + if native_format is None: + return None + schema, name, description = (native_format.get(key) for key in ("schema", "name", "description")) + return AmazonConverseConfig._create_output_config_for_response_format( + json_schema={**schema} if isinstance(schema, Mapping) else None, + name=name if isinstance(name, str) else None, + description=description if isinstance(description, str) else None, + ) + @staticmethod def _create_output_config_for_response_format( json_schema: dict | None = None, @@ -1046,6 +1029,7 @@ class AmazonConverseConfig(BaseConfig): optional_params: dict, model: str, drop_params: bool, + model_info: Mapping[str, object] | None = None, ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) base_model: Final = BedrockModelInfo.get_base_model(model) @@ -1062,6 +1046,7 @@ class AmazonConverseConfig(BaseConfig): optional_params=optional_params, non_default_params=non_default_params, is_thinking_enabled=is_thinking_enabled, + model_info=model_info, ) if param == "max_tokens" or param == "max_completion_tokens": optional_params["maxTokens"] = ( @@ -1199,7 +1184,7 @@ class AmazonConverseConfig(BaseConfig): elif param == "context_management" and isinstance(value, (dict, list)): self._map_context_management_param(value, optional_params) if param == "requestMetadata": - self._map_request_metadata_param(value, optional_params) + self._map_request_metadata_param(value, optional_params, drop_params) if param == "service_tier" and isinstance(value, str): self._map_service_tier_param(value, optional_params) @@ -1232,9 +1217,10 @@ class AmazonConverseConfig(BaseConfig): return optional_params - def _map_request_metadata_param(self, value: object, optional_params: dict) -> None: + def _map_request_metadata_param(self, value: object, optional_params: dict, drop_params: bool) -> None: if value is not None and isinstance(value, dict): - self._validate_request_metadata(value) + if not caller_request_metadata_violations_are_dropped(drop_params): + self._validate_request_metadata(value) optional_params["requestMetadata"] = value def _map_context_management_param(self, value: dict | list, optional_params: dict) -> None: @@ -1273,6 +1259,7 @@ class AmazonConverseConfig(BaseConfig): optional_params: dict, non_default_params: dict, is_thinking_enabled: bool, + model_info: Mapping[str, object] | None = None, ) -> dict: """ Handles translation of response_format parameter to Bedrock format. @@ -1296,13 +1283,23 @@ class AmazonConverseConfig(BaseConfig): if "type" in value and value["type"] == "text": return optional_params - if self._supports_native_structured_outputs(model, self.custom_llm_provider) and json_schema is not None: + native_schema: Final = ( + self._native_output_schema( + model, + json_schema, + self.is_thinking_requested(REQUEST_PARAMS.validate_python(non_default_params)), + model_info, + ) + if json_schema is not None + else None + ) + if native_schema is not None: # Use Bedrock's native structured outputs API (outputConfig.textFormat) # No synthetic tool injection, no fake_stream needed. # Requires an explicit schema — json_object with no schema falls through # to the tool-call path below. output_config: Final = self._create_output_config_for_response_format( - json_schema=json_schema, + json_schema=dict(native_schema), name=name, description=description, ) @@ -1644,7 +1641,11 @@ class AmazonConverseConfig(BaseConfig): tool_choice.pop("type", None) def _prepare_request_params( - self, optional_params: dict, model: str, drop_params: bool = False + self, + optional_params: dict, + model: str, + drop_params: bool = False, + model_info: Mapping[str, object] | None = None, ) -> tuple[dict, dict, dict, OutputConfigBlock | None]: """Prepare and separate request parameters.""" # Consume the internal ``_output_config_normalized`` marker set by @@ -1677,7 +1678,7 @@ class AmazonConverseConfig(BaseConfig): # Extract requestMetadata before processing other parameters request_metadata: Final = inference_params.pop("requestMetadata", None) - if request_metadata is not None: + if request_metadata is not None and not caller_request_metadata_violations_are_dropped(drop_params): self._validate_request_metadata(request_metadata) output_config: OutputConfigBlock | None = inference_params.pop("outputConfig", None) @@ -1687,17 +1688,12 @@ class AmazonConverseConfig(BaseConfig): and output_config_format is not None and output_config_format.get("type") == "json_schema" and base_model.startswith("anthropic") - and self._supports_native_structured_outputs(model, self.custom_llm_provider) ): - output_config = self._create_output_config_for_response_format( - json_schema=output_config_format.get("schema"), - name=output_config_format.get("name"), - description=output_config_format.get("description"), - ) - elif output_config is None and output_config_format is not None: + output_config = self._output_config_for_anthropic_format(model, output_config_format, model_info) + if output_config is None and output_config_format is not None: litellm.verbose_logger.warning( "Bedrock Converse: dropping `output_config.format` for model=%s — " - "model does not advertise `supports_native_structured_output` in " + "model does not take native structured output for this schema per " "model_prices_and_context_window.json. The schema will not be " "enforced; pass `response_format` to use the synthetic tool-call " "fallback.", @@ -1987,7 +1983,7 @@ class AmazonConverseConfig(BaseConfig): additional_request_params, request_metadata, output_config, - ) = self._prepare_request_params(optional_params, model, drop_params) + ) = self._prepare_request_params(optional_params, model, drop_params, deployment_model_info(litellm_params)) original_tools: Final = inference_params.pop("tools", []) @@ -2109,7 +2105,7 @@ class AmazonConverseConfig(BaseConfig): request_metadata: Final = resolve_bedrock_request_metadata( litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") ) - if bedrock_request_metadata_is_owned(): + if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True): _data.pop("requestMetadata", None) if request_metadata is not None: _data["requestMetadata"] = request_metadata @@ -2189,7 +2185,7 @@ class AmazonConverseConfig(BaseConfig): request_metadata: Final = resolve_bedrock_request_metadata( litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") ) - if bedrock_request_metadata_is_owned(): + if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True): _data.pop("requestMetadata", None) if request_metadata is not None: _data["requestMetadata"] = request_metadata diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index 8d1ff1d2bd9..ce28e579dcf 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -92,8 +92,16 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): optional_params: dict, model: str, drop_params: bool, + model_info: Mapping[str, object] | None = None, ) -> dict: - return AmazonConverseConfig.map_openai_params(self, non_default_params, optional_params, model, drop_params) + return AmazonConverseConfig.map_openai_params( + self, + {key: value for key, value in non_default_params.items() if key != "requestMetadata"}, + optional_params, + model, + drop_params, + model_info, + ) def transform_request( self, @@ -107,7 +115,9 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): self, model=model, messages=messages, - optional_params=_without_tool_config_injection_points(optional_params), + optional_params=_without_tool_config_injection_points( + {key: value for key, value in optional_params.items() if key != "requestMetadata"} + ), litellm_params=litellm_params, headers=headers, ) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index fc6088ddc66..ba0bb607975 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -23,7 +23,10 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation ) from litellm.llms.bedrock.common_utils import ( apply_bedrock_invoke_structured_output, + bedrock_chat_uses_native_structured_output, + bedrock_cleaned_output_schema, bedrock_supports_tool_search, + deployment_model_info, get_anthropic_beta_from_headers, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, @@ -49,6 +52,7 @@ else: LiteLLMLoggingObj = Any _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): @@ -71,6 +75,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_version: str = "bedrock-2023-05-31" + def __init__(self, model_info: Mapping[str, object] | None = None) -> None: + super().__init__() + self._deployment_model_info: Final = model_info + @property def custom_llm_provider(self) -> str | None: return "bedrock" @@ -81,10 +89,18 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def get_supported_openai_params(self, model: str) -> list[str]: return AnthropicConfig.get_supported_openai_params(self, model) - def uses_native_structured_output(self, model: str) -> bool: - from litellm.utils import supports_native_structured_output + def native_output_format_schema(self, json_schema: Mapping[str, object]) -> Mapping[str, object] | None: + return AnthropicConfig.native_output_schema(json_schema) - return supports_native_structured_output(model=model, custom_llm_provider="bedrock") + def uses_native_structured_output_for_schema( + self, + model: str, + json_schema: Mapping[str, object] | None, + request_params: Mapping[str, object], + ) -> bool: + return bedrock_chat_uses_native_structured_output( + model, self.is_thinking_requested(request_params), model_info=self._deployment_model_info + ) and (json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None) def map_openai_params( self, @@ -240,6 +256,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): apply_bedrock_invoke_structured_output( model=model, request_body=anthropic_request, + model_info=deployment_model_info(_LITELLM_PARAMS.validate_python(litellm_params)), ) strip_unsupported_bedrock_invoke_output_config_keys( model=model, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5862b05a619..17a8192ca81 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -29,6 +29,7 @@ import litellm from litellm import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.optional_imports import ensure_optional_import +from litellm.litellm_core_utils.prompt_templates.common_utils import schema_property_count from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -249,15 +250,99 @@ def convert_bedrock_invoke_output_format_to_inline_schema( request_body["messages"] = new_messages -def _bedrock_model_supports(model: str, key: str) -> bool: +def _bedrock_model_supports(model: str, key: str, custom_llm_provider: str = "bedrock") -> bool: from litellm.utils import supports_factory - return supports_factory(model=model, custom_llm_provider="bedrock", key=key) + return supports_factory(model=model, custom_llm_provider=custom_llm_provider, key=key) + + +_DEPLOYMENT_MODEL_INFO: Final = TypeAdapter(dict[str, object]) + + +def deployment_model_info(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]: + model_info: Final = (litellm_params or {}).get("model_info") + return _DEPLOYMENT_MODEL_INFO.validate_python(model_info) if isinstance(model_info, Mapping) else {} + + +def _bedrock_native_output_opt_in( + model: str, custom_llm_provider: str, model_info: Mapping[str, object] | None +) -> bool: + deployment_flag: Final = (model_info or {}).get("supports_bedrock_native_structured_output") + if isinstance(deployment_flag, bool): + return deployment_flag + return _bedrock_model_supports(model, "supports_bedrock_native_structured_output", custom_llm_provider) + + +def bedrock_supports_native_structured_output( + model: str, custom_llm_provider: str = "bedrock", model_info: Mapping[str, object] | None = None +) -> bool: + return _bedrock_model_supports( + model, "supports_native_structured_output", custom_llm_provider + ) or _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) + + +def bedrock_chat_uses_native_structured_output( + model: str, + thinking_requested: bool, + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> bool: + return _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider) or ( + thinking_requested and _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) + ) + + +BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES: Final = 64 + + +def bedrock_cleaned_output_schema( + model: str, schema: Mapping[str, object], custom_llm_provider: str = "bedrock" +) -> Mapping[str, object] | None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + if ( + not _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider) + and schema_property_count(schema) > BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES + ): + return None + return AnthropicConfig.native_output_schema(schema) + + +def bedrock_native_output_schema( + model: str, + schema: Mapping[str, object], + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> Mapping[str, object] | None: + native_schema: Final = bedrock_cleaned_output_schema(model, schema, custom_llm_provider) + if _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider): + return schema if native_schema is None else native_schema + if _bedrock_native_output_opt_in(model, custom_llm_provider, model_info): + return native_schema + return None + + +def bedrock_native_output_format( + model: str, + output_format: Mapping[str, object], + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> Mapping[str, object] | None: + schema: Final = output_format.get("schema") + if not isinstance(schema, Mapping): + return ( + output_format if bedrock_supports_native_structured_output(model, custom_llm_provider, model_info) else None + ) + native_schema: Final = bedrock_native_output_schema(model, schema, custom_llm_provider, model_info) + if native_schema is None: + return None + return {**output_format, "schema": native_schema} def apply_bedrock_invoke_structured_output( model: str, request_body: dict[str, object], # mutable-ok: edited in place like siblings + model_info: Mapping[str, object] | None = None, ) -> None: """ Route Anthropic structured-output params to what the Bedrock model supports. @@ -277,17 +362,18 @@ def apply_bedrock_invoke_structured_output( if schema_format is None: return - if _bedrock_model_supports(model, "supports_native_structured_output"): + native_format: Final = bedrock_native_output_format(model, schema_format, model_info=model_info) + if native_format is not None: existing_output_config: Final = request_body.get("output_config") if isinstance(existing_output_config, dict): - existing_output_config["format"] = schema_format + existing_output_config["format"] = native_format else: - request_body["output_config"] = {"format": schema_format} # rebind-ok: out-param + request_body["output_config"] = {"format": native_format} # rebind-ok: out-param return verbose_logger.warning( - "Bedrock Invoke: model=%s does not advertise `supports_native_structured_output` " - "in model_prices_and_context_window.json, so the JSON schema was inlined into " + "Bedrock Invoke: model=%s does not take native structured output for this schema " + "per model_prices_and_context_window.json, so the JSON schema was inlined into " "the last user message and is NOT enforced by the model.", model, ) @@ -1296,7 +1382,6 @@ def is_claude_4_5_on_bedrock(model: str) -> bool: _BEDROCK_MODEL_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") -_DEPLOYMENT_MODEL_INFO: Final = TypeAdapter(dict[str, object]) def bedrock_converse_supports_strict_tools(model: str) -> bool: @@ -1327,8 +1412,7 @@ def bedrock_model_supports_regex_lookaround(model: str, litellm_params: Mapping[ entry of ``model`` itself. A model nobody flagged keeps its schema as sent. """ params: Final = litellm_params or {} - model_info: Final = _DEPLOYMENT_MODEL_INFO.validate_python(params.get("model_info") or {}) - deployment_flag: Final = model_info.get("supports_regex_lookaround") + deployment_flag: Final = deployment_model_info(params).get("supports_regex_lookaround") if isinstance(deployment_flag, bool): return deployment_flag base_model: Final = params.get("base_model") diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4d40e0dec85..61e084fe544 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1021,7 +1021,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): converse_config: Final = AmazonConverseConfig() mapped_params = converse_config.map_openai_params( - non_default_params=optional_params, + non_default_params={k: v for k, v in optional_params.items() if k != "requestMetadata"}, optional_params={}, model=model, drop_params=False, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 95e81fa3ee7..96a7186b136 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -3,6 +3,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -69,6 +70,9 @@ else: LiteLLMLoggingObj = Any +_MODEL_INFO: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + class AmazonAnthropicClaudeMessagesConfig( AnthropicMessagesConfig, AmazonInvokeConfig, @@ -751,6 +755,7 @@ class AmazonAnthropicClaudeMessagesConfig( apply_bedrock_invoke_structured_output( model=model, request_body=anthropic_messages_request, + model_info=_MODEL_INFO.validate_python(getattr(litellm_params, "model_info", None)), ) normalize_bedrock_opus_output_config_effort( model=model, diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 912936a8ecd..ded3ce2415f 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -1,18 +1,32 @@ -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias +import base64 +import copy +import json +import struct +from binascii import crc32 +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias, TypeVar +from pydantic import TypeAdapter, ValidationError 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.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, + StreamTransformSink, +) from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, ) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.pass_through.guardrail_translation.handler import ( PassThroughEndpointHandler, @@ -21,6 +35,7 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) +_INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream" _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream" _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream" @@ -30,6 +45,10 @@ def _is_converse_endpoint(endpoint: str) -> bool: return bool(parts) and parts[-1] in _CONVERSE_ACTIONS +def _is_invoke_stream_endpoint(endpoint: str) -> bool: + return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION + + def _generic_passthrough_handler() -> "PassThroughEndpointHandler": """ Fallback for non-Converse Bedrock routes (e.g. invoke). The generic @@ -217,8 +236,6 @@ class _StreamFrame(TypedDict): def _unpack_uint32(buffer: bytes) -> int: - import struct - return struct.unpack("!I", buffer)[0] @@ -253,6 +270,578 @@ def _collect_stream_delta_text_holders(delta: object) -> list[_DeltaHolder]: return holders +def _delta_frame_texts(payload: Mapping[str, object]) -> tuple[tuple[_StreamGroupKey, str], ...]: + """The guardrailable texts of one ``contentBlockDelta`` payload, each keyed by its content block and kind.""" + block_index: Final = payload.get("contentBlockIndex") + # Group per content block: a guardrail edit that changes length must not move text between blocks. + return tuple( + ((block_index if isinstance(block_index, int) else None, group_key), container[key]) + for group_key, container, key in _collect_stream_delta_text_holders(payload.get("delta")) + ) + + +@dataclass(frozen=True, slots=True) +class _StreamTextGroups: + """Stream texts joined per group, keeping where each fragment came from so a rewrite can be split back.""" + + groups: tuple[_StreamGroupKey, ...] + members: Mapping[_StreamGroupKey, Sequence[tuple[int, int]]] + fragments: Mapping[_StreamGroupKey, Sequence[str]] + + def texts(self) -> tuple[str, ...]: + return tuple("".join(self.fragments[group]) for group in self.groups) + + +def _group_stream_texts(frame_texts: Sequence[Sequence[tuple[_StreamGroupKey, str]]]) -> _StreamTextGroups: + """Group the per-frame texts by key in first-seen order; ``groups`` keeps only the groups with any text.""" + order: Final[list[_StreamGroupKey]] = [] # mutable-ok: built in one pass over the frames + members: Final[dict[_StreamGroupKey, list[tuple[int, int]]]] = {} # mutable-ok: built in one pass over the frames + fragments: Final[dict[_StreamGroupKey, list[str]]] = {} # mutable-ok: built in one pass over the frames + for frame_idx, texts in enumerate(frame_texts): + for local_idx, (group_key, text) in enumerate(texts): + if group_key not in members: + members[group_key] = [] + fragments[group_key] = [] + order.append(group_key) + members[group_key].append((frame_idx, local_idx)) + fragments[group_key].append(text) + return _StreamTextGroups( + groups=tuple(group for group in order if "".join(fragments[group])), + members=members, + fragments=fragments, + ) + + +def _split_group_texts(groups: _StreamTextGroups, rewritten_texts: Sequence[str]) -> Mapping[tuple[int, int], str]: + """Cut each group's rewritten text back into its fragments, sized like the original fragments.""" + split: Final[dict[tuple[int, int], str]] = {} # mutable-ok: built in one pass over the groups + for group_key, rewritten in zip(groups.groups, rewritten_texts): + members = groups.members[group_key] + originals = groups.fragments[group_key] + total_orig = sum(len(t) for t in originals) or 1 + pos = 0 + for k, member in enumerate(members): + if k == len(members) - 1: + split[member] = rewritten[pos:] + else: + end = pos + round(len(rewritten) * len(originals[k]) / total_orig) + split[member] = rewritten[pos:end] + pos = end + return split + + +def _event_stream_message(headers_bytes: bytes, payload: bytes) -> bytes: + """One AWS event-stream message: prelude, prelude CRC, headers, payload and message CRC.""" + prelude: Final = struct.pack("!II", 12 + len(headers_bytes) + len(payload) + 4, len(headers_bytes)) + prelude_crc: Final = crc32(prelude) & 0xFFFFFFFF + message: Final = prelude + struct.pack("!I", prelude_crc) + headers_bytes + payload + return message + struct.pack("!I", crc32(message) & 0xFFFFFFFF) + + +def _string_header(name: str, value: str) -> bytes: + name_bytes: Final = name.encode() + value_bytes: Final = value.encode() + return bytes((len(name_bytes),)) + name_bytes + b"\x07" + struct.pack("!H", len(value_bytes)) + value_bytes + + +def _string_headers(headers: Mapping[str, str]) -> bytes: + return b"".join(_string_header(name, value) for name, value in headers.items()) + + +def _converse_event(event_type: str, payload: Mapping[str, object]) -> bytes: + return _event_stream_message( + _string_headers({":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}), + json.dumps(payload, separators=(",", ":")).encode(), + ) + + +def _converse_exception(message: str, status_code: int) -> bytes: + """An in-stream exception frame, which boto3 raises as ``EventStreamError`` carrying ``message``.""" + return _event_stream_message( + _string_headers( + { + ":exception-type": "validationException" if status_code < 500 else "internalServerException", + ":content-type": "application/json", + ":message-type": "exception", + } + ), + json.dumps({"message": message}, separators=(",", ":")).encode(), + ) + + +def _base64_json(event: Mapping[str, object]) -> str: + return base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode() + + +def _invoke_chunk(event: Mapping[str, object]) -> bytes: + return _event_stream_message( + _string_headers({":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}), + json.dumps({"bytes": _base64_json(event)}, separators=(",", ":")).encode(), + ) + + +@dataclass(slots=True) +class ConverseStreamFrame: + """One frame of a Converse event stream as the proxy relays it. + + ``payload`` is the parsed JSON of an event frame; exception frames and bytes that do not decode keep it None and + are relayed as received. ``ends_stream`` marks the empty item the framer adds once the upstream has ended, so an + exception frame or a closed connection ends the stream as ``messageStop`` does. ``guardrailed_payload`` is set + when a guardrail rewrote the frame's text.""" + + raw: bytes + event_type: str | None = None + payload: Mapping[str, object] | None = None + guardrailed_payload: Mapping[str, object] | None = None + ends_stream: bool = False + + @property + def block_index(self) -> int | None: + index: Final = self.payload.get("contentBlockIndex") if self.payload is not None else None + return index if isinstance(index, int) else None + + def scanned_texts(self) -> tuple[tuple[_StreamGroupKey, str], ...]: + current: Final = self.guardrailed_payload if self.guardrailed_payload is not None else self.payload + return _delta_frame_texts(current) if self.event_type == "contentBlockDelta" and current is not None else () + + def received_texts(self) -> tuple[tuple[_StreamGroupKey, str], ...]: + return ( + _delta_frame_texts(self.payload) + if self.event_type == "contentBlockDelta" and self.payload is not None + else () + ) + + def encode(self) -> bytes: + if self.guardrailed_payload is None: + return self.raw + headers_length: Final = _unpack_uint32(self.raw[4:8]) + return _event_stream_message( + self.raw[12 : 12 + headers_length], json.dumps(self.guardrailed_payload, separators=(",", ":")).encode() + ) + + +def _anthropic_sse(event: Mapping[str, object] | None) -> bytes: + if event is None: + return b"" + return f"event: {event['type']}\ndata: {json.dumps(event, separators=(',', ':'))}\n\n".encode() + + +@dataclass(slots=True) +class InvokeStreamFrame: + """One frame of an invoke-with-response-stream as the proxy relays it. + + ``event`` is the Anthropic Messages stream event a ``chunk`` frame carries base64 encoded in the ``bytes`` of its + ``payload``. Exception frames and bytes that do not decode keep it None and are relayed as received. + ``ends_stream`` marks the empty item the framer adds once the upstream has ended, and ``guardrailed_event`` is set + when a guardrail rewrote the event.""" + + raw: bytes + payload: Mapping[str, object] | None = None + event: Mapping[str, object] | None = None + guardrailed_event: Mapping[str, object] | None = None + ends_stream: bool = False + + def sse(self) -> bytes: + """The event as a ``/v1/messages`` stream carries it, the shape its guardrail translation scans.""" + return _anthropic_sse(self.guardrailed_event if self.guardrailed_event is not None else self.event) + + def received_sse(self) -> bytes: + """The event as received from the upstream, before any guardrail rewrote it.""" + return _anthropic_sse(self.event) + + def encode(self) -> bytes: + if self.guardrailed_event is None or self.payload is None: + return self.raw + headers_length: Final = _unpack_uint32(self.raw[4:8]) + return _event_stream_message( + self.raw[12 : 12 + headers_length], + json.dumps({**self.payload, "bytes": _base64_json(self.guardrailed_event)}, separators=(",", ":")).encode(), + ) + + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_Frame: Final = TypeVar("_Frame", ConverseStreamFrame, InvokeStreamFrame) + + +def _event_message(raw: bytes) -> tuple[str, bytes] | None: + """The event type and payload of one event-stream frame, or None for an exception frame or bytes that do not + decode.""" + from botocore.eventstream import EventStreamBuffer, ParserError + + buffer: Final = EventStreamBuffer() + buffer.add_data(raw) + try: + message: Final = next(iter(buffer)) + except (ParserError, StopIteration) as e: + verbose_proxy_logger.debug( + "BedrockPassthroughGuardrailHandler: could not decode event-stream frame, relaying it unmodified: %s", e + ) + return None + # botocore-stubs types EventStreamMessage.headers as int; at runtime it is the decoded header dict. + headers: Final = _JSON_OBJECT.validate_python(message.headers) + event_type: Final = headers.get(":event-type") + if headers.get(":message-type") != "event" or not isinstance(event_type, str): + return None + return event_type, message.payload + + +def _decode_converse_frame(raw: bytes) -> ConverseStreamFrame: + message: Final = _event_message(raw) + if message is None: + return ConverseStreamFrame(raw=raw) + event_type, payload_bytes = message + try: + payload: Final = _JSON_OBJECT.validate_python(json.loads(payload_bytes)) + except (json.JSONDecodeError, UnicodeDecodeError, ValidationError): + return ConverseStreamFrame(raw=raw, event_type=event_type) + return ConverseStreamFrame(raw=raw, event_type=event_type, payload=payload) + + +def _decode_invoke_frame(raw: bytes) -> InvokeStreamFrame: + message: Final = _event_message(raw) + if message is None or message[0] != "chunk": + return InvokeStreamFrame(raw=raw) + try: + payload: Final = _JSON_OBJECT.validate_python(json.loads(message[1])) + encoded: Final = payload.get("bytes") + event: Final = ( + _JSON_OBJECT.validate_python(json.loads(base64.b64decode(encoded))) if isinstance(encoded, str) else None + ) + except (ValueError, ValidationError): + return InvokeStreamFrame(raw=raw) + if event is None or not isinstance(event.get("type"), str): + return InvokeStreamFrame(raw=raw) + return InvokeStreamFrame(raw=raw, payload=payload, event=event) + + +async def _event_stream_frames( + chunks: AsyncIterable[bytes], frame_type: type[_Frame], decode: Callable[[bytes], _Frame] +) -> AsyncGenerator[_Frame, None]: + """Cut the upstream bytes into frames as soon as each frame has fully arrived, then add an empty item that ends + the stream once the upstream has ended. + + Bytes that cannot be framed (a length below the 16-byte minimum, or a stream that ends inside a frame) are relayed + as received.""" + pending = b"" # rebind-ok: bytes received but not yet framed + framing = True # rebind-ok: cleared once a frame length is unusable + async for chunk in chunks: + if not framing: + yield frame_type(raw=chunk) + continue + pending += chunk + while len(pending) >= 4: + total_length = _unpack_uint32(pending[:4]) + if total_length < 16: + framing = False + break + if len(pending) < total_length: + break + yield decode(pending[:total_length]) + pending = pending[total_length:] + if not framing and pending: + yield frame_type(raw=pending) + pending = b"" + if pending: + yield frame_type(raw=pending) + yield frame_type(raw=b"", ends_stream=True) + + +def _converse_frames(items: Sequence[object]) -> tuple[ConverseStreamFrame, ...]: + """The frames among ``items``, with the bytes an earlier guardrail's block message added read back into frames.""" + return tuple( + item if isinstance(item, ConverseStreamFrame) else _decode_converse_frame(item) + for item in items + if isinstance(item, (ConverseStreamFrame, bytes)) + ) + + +def _reads_stream_marks(guardrail: "CustomGuardrail") -> bool: + return getattr(guardrail, "reads_stream_marks", False) is True + + +def _json_object_member(value: object, key: str) -> Mapping[str, object] | None: + """``value[key]`` when both are JSON objects, else None.""" + try: + return _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(value).get(key)) + except ValidationError: + return None + + +def _reject_json_constant(name: str) -> object: + raise ValueError(f"{name} is not JSON") + + +def _finished_tool_input(text: str) -> object: + """A tool call's streamed input as the JSON document Converse returns, or the text as received when it is not + JSON (cut short, rewritten by a guardrail, or holding NaN or nesting too deep to send on).""" + try: + return json.loads(text, parse_constant=_reject_json_constant) if text else {} + except (ValueError, RecursionError): + return text + + +def _finished_converse_content( + frames: Sequence[ConverseStreamFrame], groups: _StreamTextGroups +) -> list[dict[str, object]]: # mutable-ok: the content of the finished response guardrails receive + """The finished reply's text, reasoning and tool-call blocks, shaped as a non-streaming Converse response + holds them.""" + texts: Final = dict(zip(groups.groups, groups.texts())) + tool_starts: Final[dict[int | None, Mapping[str, object]]] = {} # mutable-ok: built in one pass over the frames + keys: Final[dict[_StreamGroupKey, None]] = {} # mutable-ok: blocks in stream order, built in one pass + for frame in frames: + tool_use = ( + _json_object_member(_json_object_member(frame.payload, "start"), "toolUse") + if frame.event_type == "contentBlockStart" + else None + ) + if tool_use is not None: + tool_starts[frame.block_index] = tool_use + keys[(frame.block_index, "tool")] = None + keys.update((key, None) for key, _ in frame.scanned_texts()) + content: Final[list[dict[str, object]]] = [] # mutable-ok: built in one pass over the blocks + for block_index, kind in keys: + text = texts.get((block_index, kind), "") + if kind == "tool": + tool_start = tool_starts.get(block_index, {}) + named = {name: tool_start[name] for name in ("toolUseId", "name", "type") if name in tool_start} + content.append({"toolUse": {**named, "input": _finished_tool_input(text)}}) + elif text and kind == "reasoning": + content.append({"reasoningContent": {"reasoningText": {"text": text}}}) + elif text: + content.append({"text": text}) + return content + + +def _stream_marked_request_data( + request_data: Mapping[str, object], frames: Sequence[ConverseStreamFrame], groups: _StreamTextGroups, ended: bool +) -> dict[str, object]: # mutable-ok: the request dict guardrails receive + """``request_data`` with the marks chat completions gives a streamed answer: ``stream`` on every scan and, once + the stream has ended, the finished ``response``.""" + if not ended: + return {**request_data, "stream": True} + stop_reason: Final = next( + ( + frame.payload.get("stopReason") + for frame in frames + if frame.event_type == "messageStop" and frame.payload is not None + ), + None, + ) + message: Final = {"role": "assistant", "content": _finished_converse_content(frames, groups)} + response: Final = {"output": {"message": message}, **({"stopReason": stop_reason} if stop_reason else {})} + return {**request_data, "stream": True, "response": response} + + +def _invoke_frames(items: Sequence[object]) -> tuple[InvokeStreamFrame, ...]: + return tuple(item for item in items if isinstance(item, InvokeStreamFrame)) + + +def _sse_events(sse: bytes) -> tuple[Mapping[str, object], ...]: + return tuple( + _JSON_OBJECT.validate_python(json.loads(line[len("data:") :])) + for line in sse.decode().splitlines() + if line.startswith("data:") + ) + + +def _invoke_message_stopped(frames: Sequence[InvokeStreamFrame]) -> bool: + """Whether a ``message_delta`` with a ``stop_reason`` came, as a ``/v1/messages`` stream ends.""" + return any( + frame.event is not None + and frame.event.get("type") == "message_delta" + and isinstance(delta := frame.event.get("delta"), Mapping) + and _JSON_OBJECT.validate_python(delta).get("stop_reason") is not None + for frame in frames + ) + + +def _invoke_stream_has_ended(frames: Sequence[InvokeStreamFrame]) -> bool: + """Ended on a ``message_delta`` with a ``stop_reason`` or once the upstream has ended.""" + return any(frame.ends_stream for frame in frames) or _invoke_message_stopped(frames) + + +_UPSTREAM_ENDED_SSE: Final = _anthropic_sse( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 0}, + } +) + + +def _invoke_scan_sse(frames: Sequence[InvokeStreamFrame], events: Sequence[bytes]) -> tuple[bytes, ...]: + """The ``/v1/messages`` events to scan for these frames, closed with a ``message_delta`` when the upstream ended + without one so the translation scans them as a whole stream, tool calls included.""" + if any(frame.ends_stream for frame in frames) and not _invoke_message_stopped(frames): + return (*events, _UPSTREAM_ENDED_SSE) + return tuple(events) + + +async def _scan_invoke_stream( + frames: Sequence[InvokeStreamFrame], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"], + user_api_key_dict: Optional["UserAPIKeyAuth"], + request_data: dict[str, object] | None, # mutable-ok: the request dict guardrails receive +) -> None: + """Scan the Anthropic events relayed so far as ``/v1/messages`` scans its stream, with the request marked as + streaming and those events as its ``responses``. Once the stream has ended, the rewrites are written back into the + frames, which reach the caller only when the stream is held until moderated.""" + received: Final = tuple(frame.sse() for frame in frames) + scan_events: Final = _invoke_scan_sse(frames, received) + scanned: Final[list[object]] = list(scan_events) # mutable-ok: the /v1/messages translation rewrites it in place + if request_data is not None: + request_data["stream"] = True # rebind-ok: marks the request guardrails receive as a stream + request_data["responses"] = scanned # rebind-ok: each scan sees every event relayed so far + await AnthropicMessagesHandler().process_output_streaming_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + responses_so_far=scanned, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + deliver_ended_stream_rewrites=_invoke_stream_has_ended(frames), + ) + for frame, before, after in zip(frames, received, scanned): + if isinstance(after, bytes) and after != before and (rewritten := _sse_events(after)): + frame.guardrailed_event = rewritten[0] + + +def _invoke_block_frames( + exc: "ModifyResponseException", stream_started: bool, frames: Sequence[InvokeStreamFrame] +) -> tuple[bytes, ...]: + """The ``/v1/messages`` block events for the stream relayed so far, each in its own ``chunk`` frame.""" + sse_chunks: Final = AnthropicMessagesHandler().build_block_sse_chunks( + exc, stream_started=stream_started, responses_so_far=[frame.sse() for frame in frames] + ) + return tuple(_invoke_chunk(event) for event in _sse_events(b"".join(sse_chunks))) + + +def _write_back_stream_texts( + frames: Sequence[ConverseStreamFrame], groups: _StreamTextGroups, rewritten_texts: Sequence[str] +) -> None: + split: Final = _split_group_texts(groups, rewritten_texts) + for frame_idx, frame in enumerate(frames): + current = frame.guardrailed_payload if frame.guardrailed_payload is not None else frame.payload + if frame.event_type != "contentBlockDelta" or current is None: + continue + rewritten = copy.deepcopy(current) + changed = False + for local_idx, (_, container, key) in enumerate(_collect_stream_delta_text_holders(rewritten.get("delta"))): + new_text = split.get((frame_idx, local_idx)) + if new_text is not None and new_text != container[key]: + container[key] = new_text + changed = True + if changed: + frame.guardrailed_payload = rewritten + + +_INVOKE_DELTA_TEXT_FIELDS: Final = MappingProxyType( + { + "text_delta": ("text", "text"), + "input_json_delta": ("tool", "partial_json"), + "thinking_delta": ("reasoning", "thinking"), + } +) + + +@dataclass(frozen=True, slots=True) +class _InvokeDeltaText: + """The guardrailable text one Anthropic event's delta carries, keyed by its content block and kind like the + Converse holders: text, tool input and thinking. Signatures and citations are never rewritten.""" + + group: _StreamGroupKey + delta: Mapping[str, object] + field: str + text: str + + +def _invoke_delta_text(event: Mapping[str, object] | None) -> _InvokeDeltaText | None: + if event is None: + return None + try: + delta: Final = _JSON_OBJECT.validate_python(event.get("delta")) + except ValidationError: + return None + kind: Final = delta.get("type") + field: Final = _INVOKE_DELTA_TEXT_FIELDS.get(kind) if isinstance(kind, str) else None + text: Final = delta.get(field[1]) if field is not None else None + if field is None or not isinstance(text, str): + return None + index: Final = event.get("index") + return _InvokeDeltaText((index if isinstance(index, int) else None, field[0]), delta, field[1], text) + + +async def _post_call_rewritten_texts( + groups: _StreamTextGroups, + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: "UserAPIKeyAuth", + data: dict[str, object], # mutable-ok: the request dict post_call hooks receive +) -> Sequence[str] | None: + """Run the post_call guardrails on a buffered stream's texts as one Converse response with a content block per + group, and return each group's text as the guardrails left it, or None when the response came back in another + shape.""" + synthetic_response: Final[dict] = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": text} for text in groups.texts()], + } + }, + "stopReason": "end_turn", + } + + processed: Final = await proxy_logging_obj.post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=synthetic_response, + ) + + if not isinstance(processed, dict): + verbose_proxy_logger.debug( + "BedrockPassthroughGuardrailHandler: post_call_success_hook returned %s, leaving event stream unmodified", + type(processed).__name__, + ) + return None + + try: + processed_blocks: Final = processed["output"]["message"]["content"] + return [processed_blocks[i]["text"] for i in range(len(groups.groups))] + except (KeyError, IndexError, TypeError): + return None + + +async def _guardrail_buffered_invoke_stream( + body_bytes: bytes, + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: "UserAPIKeyAuth", + data: dict[str, object], # mutable-ok: the request dict post_call hooks receive +) -> bytes: + """Scan a whole Claude invoke stream the way the buffered Converse path scans its stream, and write the guardrail + rewrites back into the ``chunk`` frames that carried each text.""" + + async def received() -> AsyncGenerator[bytes, None]: + yield body_bytes + + frames: Final = [frame async for frame in _event_stream_frames(received(), InvokeStreamFrame, _decode_invoke_frame)] + texts: Final = [_invoke_delta_text(frame.event) for frame in frames] + groups: Final = _group_stream_texts([[] if text is None else [(text.group, text.text)] for text in texts]) + if not groups.groups: + return body_bytes + rewritten_texts: Final = await _post_call_rewritten_texts(groups, proxy_logging_obj, user_api_key_dict, data) + if rewritten_texts is None: + return body_bytes + split: Final = _split_group_texts(groups, rewritten_texts) + for frame_idx, (frame, text) in enumerate(zip(frames, texts)): + if ( + text is not None + and frame.event is not None + and (new_text := split.get((frame_idx, 0))) is not None + and new_text != text.text + ): + frame.guardrailed_event = {**frame.event, "delta": {**text.delta, text.field: new_text}} + return b"".join(frame.encode() for frame in frames) + + class BedrockPassthroughGuardrailHandler(BaseTranslation): @staticmethod def is_event_stream_content_type(content_type: str) -> bool: @@ -263,19 +852,32 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): return _EVENT_STREAM_MEDIA_TYPE @staticmethod - def event_stream_endpoint_is_de_anonymizable(endpoint: str) -> bool: - return _is_converse_endpoint(endpoint) + def event_stream_endpoint_is_de_anonymizable(endpoint: str, model: str, body: object) -> bool: + """The streams the buffered path can scan whole and rewrite, the same ones the relay scans.""" + return BedrockPassthroughGuardrailHandler.event_stream_endpoint_is_relayable(endpoint, model, body) + + @staticmethod + def event_stream_endpoint_is_relayable(endpoint: str, model: str, body: object) -> bool: + """Converse streams of every model, and invoke streams of Claude, whose chunks carry Anthropic Messages stream + events. Claude is told by its model id or, for ids that do not name it such as application inference profile + ARNs, by the ``anthropic_version`` Bedrock requires in every Claude Messages request body.""" + return _is_converse_endpoint(endpoint) or ( + _is_invoke_stream_endpoint(endpoint) + and ( + BaseAWSLLM.get_bedrock_invoke_provider(model) == "anthropic" + or (isinstance(body, Mapping) and "anthropic_version" in body) + ) + ) @staticmethod async def de_anonymize_event_stream( body_bytes: bytes, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", - data: dict, + data: dict[str, object], # mutable-ok: the request dict post_call hooks receive ) -> bytes: - import json as _json - import struct - from binascii import crc32 as esm_crc32 + if _is_invoke_stream_endpoint(str(data.get("endpoint") or "")): + return await _guardrail_buffered_invoke_stream(body_bytes, proxy_logging_obj, user_api_key_dict, data) from botocore.eventstream import EventStreamBuffer @@ -304,16 +906,11 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): frames.append({"raw": frame_raw, "texts": []}) continue - texts: list[tuple[_StreamGroupKey, str]] = [] + texts: Sequence[tuple[_StreamGroupKey, str]] = () if event_type == "contentBlockDelta": try: - payload_dict: dict[str, object] = _json.loads(payload_bytes) - block_index = payload_dict.get("contentBlockIndex") - # Group per content block: a guardrail edit that changes length must not move text between blocks. - texts = [ - ((block_index if isinstance(block_index, int) else None, group_key), container[key]) - for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta")) - ] + payload_dict: dict[str, object] = json.loads(payload_bytes) + texts = _delta_frame_texts(payload_dict) except Exception as e: verbose_proxy_logger.debug( "BedrockPassthroughGuardrailHandler: could not parse " @@ -325,66 +922,17 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): trailing_bytes: Final = body_bytes[offset:] - group_order: Final[list[_StreamGroupKey]] = [] - group_members: Final[dict[_StreamGroupKey, list[tuple[int, int]]]] = {} - group_texts: Final[dict[_StreamGroupKey, list[str]]] = {} - for frame_idx, frame in enumerate(frames): - for local_idx, (group_key, text) in enumerate(frame["texts"]): - if group_key not in group_members: - group_members[group_key] = [] - group_texts[group_key] = [] - group_order.append(group_key) - group_members[group_key].append((frame_idx, local_idx)) - group_texts[group_key].append(text) - - active_groups: Final = [gk for gk in group_order if "".join(group_texts[gk])] - if not active_groups: + groups: Final = _group_stream_texts([frame["texts"] for frame in frames]) + if not groups.groups: return body_bytes - synthetic_response: Final[dict] = { - "output": { - "message": { - "role": "assistant", - "content": [{"text": "".join(group_texts[gk])} for gk in active_groups], - } - }, - "stopReason": "end_turn", - } - - processed: Final = await proxy_logging_obj.post_call_success_hook( - data=data, - user_api_key_dict=user_api_key_dict, - response=synthetic_response, + de_anonymized_texts: Final = await _post_call_rewritten_texts( + groups, proxy_logging_obj, user_api_key_dict, data ) - - if not isinstance(processed, dict): - verbose_proxy_logger.debug( - "BedrockPassthroughGuardrailHandler: post_call_success_hook returned %s, " - "leaving event stream unmodified", - type(processed).__name__, - ) + if de_anonymized_texts is None: return body_bytes - try: - processed_blocks: Final = processed["output"]["message"]["content"] - de_anonymized_texts: Final = [processed_blocks[i]["text"] for i in range(len(active_groups))] - except (KeyError, IndexError, TypeError): - return body_bytes - - new_text_map: Final[dict[tuple[int, int], str]] = {} - for group_key, de_anonymized_text in zip(active_groups, de_anonymized_texts): - members = group_members[group_key] - orig_texts = group_texts[group_key] - total_orig = sum(len(t) for t in orig_texts) or 1 - de_anon_len = len(de_anonymized_text) - pos = 0 - for k, member in enumerate(members): - if k == len(members) - 1: - new_text_map[member] = de_anonymized_text[pos:] - else: - end = pos + round(de_anon_len * len(orig_texts[k]) / total_orig) - new_text_map[member] = de_anonymized_text[pos:end] - pos = end + new_text_map: Final = _split_group_texts(groups, de_anonymized_texts) result_parts: Final[list[bytes]] = [] @@ -399,27 +947,19 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): headers_bytes = frame_raw[12 : 12 + orig_hdrs_len] try: - payload_dict = _json.loads(frame_raw[12 + orig_hdrs_len : orig_total - 4]) + payload_dict = json.loads(frame_raw[12 + orig_hdrs_len : orig_total - 4]) for local_idx, (_, container, key) in enumerate( _collect_stream_delta_text_holders(payload_dict.get("delta")) ): new_text = new_text_map.get((frame_idx, local_idx)) if new_text is not None: container[key] = new_text - new_payload = _json.dumps(payload_dict, separators=(",", ":")).encode() + new_payload = json.dumps(payload_dict, separators=(",", ":")).encode() except Exception: result_parts.append(frame_raw) continue - new_total = 12 + orig_hdrs_len + len(new_payload) + 4 - prelude = struct.pack("!II", new_total, orig_hdrs_len) - prelude_crc_val = esm_crc32(prelude) & 0xFFFFFFFF - prelude_crc_b = struct.pack("!I", prelude_crc_val) - part_for_msg_crc = prelude_crc_b + headers_bytes + new_payload - msg_crc_val = esm_crc32(part_for_msg_crc, prelude_crc_val) & 0xFFFFFFFF - msg_crc_b = struct.pack("!I", msg_crc_val) - - result_parts.append(prelude + prelude_crc_b + headers_bytes + new_payload + msg_crc_b) + result_parts.append(_event_stream_message(headers_bytes, new_payload)) result_parts.append(trailing_bytes) return b"".join(result_parts) @@ -477,8 +1017,13 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, ) -> object: - endpoint: Final = (request_data or {}).get("endpoint", "") - if endpoint and not _is_converse_endpoint(endpoint): + endpoint: Final[object] = (request_data or {}).get("endpoint", "") + if ( + isinstance(endpoint, str) + and endpoint + and not _is_converse_endpoint(endpoint) + and not _is_invoke_stream_endpoint(endpoint) + ): return await _generic_passthrough_handler().process_output_response( response=response, guardrail_to_apply=guardrail_to_apply, @@ -503,6 +1048,20 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): if not texts: return response + await self._scan_output_texts( + texts, holders, guardrail_to_apply, litellm_logging_obj, user_api_key_dict, request_data + ) + return response + + async def _scan_output_texts( + self, + texts: list[str], + holders: list[_StringHolder], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"], + user_api_key_dict: Optional["UserAPIKeyAuth"], + request_data: dict[str, object] | None, # mutable-ok: the request dict guardrails receive + ) -> None: effective_request_data = request_data or {} if "litellm_metadata" not in effective_request_data and user_api_key_dict is not None: user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) @@ -513,8 +1072,8 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): } inputs: Final = GenericGuardrailAPIInputs(texts=texts) - model: Final = effective_request_data.get("model") if effective_request_data else None - if model: + model: Final = effective_request_data.get("model") + if isinstance(model, str) and model: inputs["model"] = model guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( @@ -528,4 +1087,141 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): if guardrailed_texts: _write_back_texts(guardrailed_texts, holders) - return response + async def process_output_streaming_response( + self, + responses_so_far: Sequence[object], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: dict[str, object] | None = None, # mutable-ok: the request dict guardrails receive + stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, + ) -> Sequence[object]: + """ + Scan the Converse frames relayed so far as the buffered path scans a whole response: one text per content + block and kind. Once the stream has ended, guardrail rewrites are written back into the frames, which reach + the caller only when the stream is held until moderated. A guardrail that sets ``reads_stream_marks`` gets + a copy of the request marked as chat completions marks a streamed answer; the others get it unchanged. + Such a guardrail also gets the ended scan of a reply with no text, such as a call to a tool that takes no + input, with no texts to scan. Invoke frames are scanned as Anthropic Messages events. + """ + invoke_frames: Final = _invoke_frames(responses_so_far) + if invoke_frames: + await _scan_invoke_stream( + invoke_frames, guardrail_to_apply, litellm_logging_obj, user_api_key_dict, request_data + ) + return responses_so_far + frames: Final = _converse_frames(responses_so_far) + groups: Final = _group_stream_texts([frame.scanned_texts() for frame in frames]) + ended: Final = self._check_streaming_has_ended(frames) + reads_marks: Final = _reads_stream_marks(guardrail_to_apply) + if not groups.groups: + if request_data is not None and reads_marks and ended and _finished_converse_content(frames, groups): + await self._scan_output_texts( + [], + [], + guardrail_to_apply, + litellm_logging_obj, + user_api_key_dict, + _stream_marked_request_data(request_data, frames, groups, ended), + ) + return responses_so_far + content: Final = [{"text": text} for text in groups.texts()] + await self.process_output_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + response={"output": {"message": {"role": "assistant", "content": content}}}, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=( + _stream_marked_request_data(request_data, frames, groups, ended) + if request_data is not None and reads_marks + else request_data + ), + ) + if ended: + _write_back_stream_texts(frames, groups, [block["text"] for block in content]) + return responses_so_far + + def needs_ended_stream_scan(self, guardrail: "CustomGuardrail", responses_so_far: Sequence[object]) -> bool: + return _reads_stream_marks(guardrail) + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + invoke_frames: Final = _invoke_frames(responses_so_far) + if invoke_frames: + return AnthropicMessagesHandler().get_streaming_scan_key( + list(_invoke_scan_sse(invoke_frames, [frame.received_sse() for frame in invoke_frames])) + ) + frames: Final = _converse_frames(responses_so_far) + return StreamingScanKey( + texts=_group_stream_texts([frame.received_texts() for frame in frames]).texts(), + stream_ended=self._check_streaming_has_ended(frames), + ) + + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: + return _invoke_stream_has_ended(_invoke_frames(responses_so_far)) or any( + frame.event_type == "messageStop" or frame.ends_stream for frame in _converse_frames(responses_so_far) + ) + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + """ + End the message with the guardrail's message as a new text block and ``stopReason`` ``guardrail_intervened``, + closing any content block left open and starting the message when nothing was relayed yet. An invoke stream + gets the ``/v1/messages`` block events instead. + """ + invoke_frames: Final = _invoke_frames(responses_so_far or ()) + if invoke_frames or _is_invoke_stream_endpoint(str(exc.request_data.get("endpoint") or "")): + return _invoke_block_frames(exc, stream_started, invoke_frames) + frames: Final = _converse_frames(responses_so_far or ()) + started: Final = frozenset( + index + for index in ( + frame.block_index for frame in frames if frame.event_type in ("contentBlockStart", "contentBlockDelta") + ) + if index is not None + ) + stopped: Final = frozenset( + index + for index in (frame.block_index for frame in frames if frame.event_type == "contentBlockStop") + if index is not None + ) + block_index: Final = max(started | stopped, default=-1) + 1 + return ( + *( + () + if any(frame.event_type == "messageStart" for frame in frames) + else (_converse_event("messageStart", {"role": "assistant"}),) + ), + *(_converse_event("contentBlockStop", {"contentBlockIndex": index}) for index in sorted(started - stopped)), + _converse_event("contentBlockDelta", {"contentBlockIndex": block_index, "delta": {"text": exc.message}}), + _converse_event("contentBlockStop", {"contentBlockIndex": block_index}), + _converse_event("messageStop", {"stopReason": "guardrail_intervened"}), + ) + + @staticmethod + def owns_event_stream_item(item: object) -> bool: + return isinstance(item, (ConverseStreamFrame, InvokeStreamFrame)) + + @staticmethod + def event_stream_items( + chunks: AsyncIterable[bytes], endpoint: str + ) -> AsyncGenerator[ConverseStreamFrame | InvokeStreamFrame, None]: + if _is_invoke_stream_endpoint(endpoint): + return _event_stream_frames(chunks, InvokeStreamFrame, _decode_invoke_frame) + return _event_stream_frames(chunks, ConverseStreamFrame, _decode_converse_frame) + + @staticmethod + def encode_event_stream_item(item: object) -> bytes: + if isinstance(item, (ConverseStreamFrame, InvokeStreamFrame)): + return item.encode() + if isinstance(item, bytes): + return item + raise TypeError(f"cannot relay {type(item).__name__} on a Bedrock event stream") + + @staticmethod + def event_stream_error_item(message: str, status_code: int) -> bytes: + return _converse_exception(message, status_code) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 03e2ee6745c..36cd87aff06 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -1,18 +1,33 @@ import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence +from functools import partial +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, cast import httpx from httpx import Response +from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.constants import AWS_ERROR_TYPE_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector +from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo +from ..request_metadata import ( + BEDROCK_REQUEST_METADATA_HEADER, + BedrockRequestMetadataViolation, + bedrock_request_metadata_headers, + bedrock_request_metadata_is_owned, + bedrock_request_metadata_violation, + caller_request_metadata_violations_are_dropped, + merge_bedrock_invoke_headers, + resolve_bedrock_request_metadata, +) if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage @@ -24,10 +39,87 @@ if TYPE_CHECKING: BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"}) +_BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) +_BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) + + +def _bedrock_endpoint_action(endpoint: str) -> str: + return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] def is_bedrock_streaming_endpoint(endpoint: str) -> bool: - return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS + return _bedrock_endpoint_action(endpoint) in BEDROCK_STREAMING_ACTIONS + + +def _caller_request_metadata_violation( + action: str, body: Mapping[str, object] +) -> BedrockRequestMetadataViolation | None: + if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body and body["requestMetadata"] is None: + return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary") + return bedrock_request_metadata_violation(body.get("requestMetadata")) + + +def _aws_error_type(action: str, violation: BedrockRequestMetadataViolation) -> str: + if action in _BEDROCK_CONVERSE_ACTIONS and violation.kind == "type": + return "SerializationException" + return "ValidationException" + + +def _caller_request_metadata(action: str, body: Mapping[str, object], litellm_params: Mapping[str, object]) -> object: + violation: Final = ( + None + if caller_request_metadata_violations_are_dropped(litellm_params.get("drop_params") is True) + else _caller_request_metadata_violation(action, body) + ) + if violation is not None: + raise BedrockError( + status_code=400, + message=violation.message, + headers={AWS_ERROR_TYPE_HEADER: _aws_error_type(action, violation)}, + ) + return body.get("requestMetadata") + + +def _with_body_request_metadata( + action: str, body: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + resolved: Final = resolve_bedrock_request_metadata( + litellm_params, caller_metadata=_caller_request_metadata(action, body, litellm_params) + ) + without_caller: Final = {key: value for key, value in body.items() if key != "requestMetadata"} + return without_caller if resolved is None else {**without_caller, "requestMetadata": resolved} + + +def _without_owned_request_metadata_header(headers: Mapping[str, str]) -> Mapping[str, str]: + if not bedrock_request_metadata_is_owned(): + return headers + return {name: value for name, value in headers.items() if name.lower() != BEDROCK_REQUEST_METADATA_HEADER.lower()} + + +def _with_request_metadata( + action: str | None, + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if action in _BEDROCK_CONVERSE_ACTIONS: + return headers, _with_body_request_metadata(action, body, litellm_params) + if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body: + return _without_owned_request_metadata_header(headers), _with_body_request_metadata( + action, body, litellm_params + ) + if action in _BEDROCK_INVOKE_ACTIONS: + return _with_invoke_request_metadata(headers, body, litellm_params) + return headers, body + + +def _with_invoke_request_metadata( + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> tuple[Mapping[str, str], Mapping[str, object]]: + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) @@ -91,6 +183,57 @@ def _translate_message(decoder: "AWSEventStreamDecoder", message: str) -> ModelR return None +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_JSON_OBJECTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def _leaf_texts(block: Mapping[str, object]) -> Iterator[str]: + """The prompt text a content block carries itself: its text or JSON, or a tool call's input.""" + text: Final = block.get("text") + if isinstance(text, str): + yield text + if "json" in block: + yield json.dumps(block["json"], default=str) + if "toolUse" in block: + yield json.dumps(_JSON_OBJECT.validate_python(block["toolUse"]).get("input"), default=str) + + +def _block_texts(block: Mapping[str, object]) -> Iterator[str]: + """The prompt text a request content block carries, including the blocks of a tool result's content, which + Bedrock does not nest further.""" + yield from _leaf_texts(block) + if "toolResult" not in block: + return + tool_result: Final = _JSON_OBJECT.validate_python(block["toolResult"]) + yield from chain.from_iterable(map(_leaf_texts, _JSON_OBJECTS.validate_python(tool_result.get("content", ())))) + + +def _content_text(content: object) -> str: + if isinstance(content, str): + return content + return "\n".join(chain.from_iterable(map(_block_texts, _JSON_OBJECTS.validate_python(content)))) + + +def _chat_message(turn: Mapping[str, object]) -> ChatCompletionUserMessage: + return ChatCompletionUserMessage(role="user", content=_content_text(turn.get("content", ""))) + + +def _count_prompt_tokens(model: str, request_body: object) -> int: + """Estimate a Bedrock request's prompt tokens from its system prompt and turns, the way a chat completions + stream that ended before the provider reported usage is estimated.""" + from litellm.litellm_core_utils.token_counter import token_counter + + body: Final = _JSON_OBJECT.validate_python(request_body) + if "messages" not in body: + return 0 + system: Final = body.get("system") + messages: Final = ( + *((ChatCompletionSystemMessage(role="system", content=_content_text(system)),) if system else ()), + *map(_chat_message, _JSON_OBJECTS.validate_python(body["messages"])), + ) + return token_counter(model=model, messages=messages) + + def _build_logged_response( chunks: Sequence[ModelResponseStream], litellm_logging_obj: "LiteLLMLoggingObj" ) -> Optional["CostResponseTypes"]: @@ -98,7 +241,13 @@ def _build_logged_response( if len(chunks) == 0: return None - return stream_chunk_builder(chunks=list(chunks), logging_obj=litellm_logging_obj) + return stream_chunk_builder( + chunks=list(chunks), + logging_obj=litellm_logging_obj, + count_prompt_tokens=partial( + _count_prompt_tokens, litellm_logging_obj.model, litellm_logging_obj.model_call_details.get("input") + ), + ) class BedrockEventStreamCollector: @@ -223,18 +372,24 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD def sign_request( self, - headers: dict, - litellm_params: dict, - request_data: dict | None, + headers: dict[str, str], + litellm_params: dict[str, object], + request_data: dict[str, object] | None, api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: optional_params: Final = litellm_params.copy() + signed_headers, signed_body = _with_request_metadata( + action=_bedrock_endpoint_action(api_base) if bedrock_request_metadata_is_owned() else None, + headers=headers, + body=request_data or {}, + litellm_params=litellm_params, + ) return self._sign_request( service_name="bedrock", - headers=headers, + headers=dict(signed_headers), optional_params=optional_params, - request_data=request_data or {}, + request_data=dict(signed_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index c0a382b136c..3e7c2f7eee9 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -8,7 +8,7 @@ import asyncio import contextlib import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, NoReturn, Protocol, runtime_checkable @@ -21,7 +21,7 @@ from litellm._uuid import uuid from litellm.constants import ( BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, - BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, @@ -37,9 +37,8 @@ from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection +from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities -_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -177,25 +176,35 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: return -def _pending_session_update(scope: Mapping[str, object]) -> str | None: - """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: - committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) - raise BedrockError( - status_code=400, - message=( - "Bedrock realtime session already committed to a provider stream; it cannot be replayed" - + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") - ), - ) - pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY) +def _is_session_committed(scope: Mapping[str, object]) -> bool: + return scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True + + +def raise_if_session_committed(scope: Mapping[str, object] | None) -> None: + """A fallback attempt on a websocket whose Bedrock stream already took input would silently start over.""" + if scope is None or not _is_session_committed(scope): + return + committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) + raise BedrockError( + status_code=400, + message=( + "Bedrock realtime session already committed to a provider stream; it cannot be replayed" + + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") + ), + ) + + +def _pending_opening_message(scope: Mapping[str, object]) -> str | None: + """A fallback attempt on the same websocket replays the opening message Bedrock never accepted a stream for.""" + raise_if_session_committed(scope) + pending: Final = scope.get(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY) return pending if isinstance(pending, str) else None def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn: error: Final = _as_bedrock_error(failure) verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", redact_string(str(error))) - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: + if _is_session_committed(scope): scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error)) raise error from failure @@ -207,26 +216,36 @@ def _parse_client_message(message: str) -> Mapping[str, JsonValue]: return _EMPTY_JSON_OBJECT -async def _ack_session_update( +async def _send_opening_message( client_ws: RealtimeClientWebSocket, bedrock_stream: BedrockBidirectionalStream, + message: str, + bedrock_messages: Sequence[str], + send_to_bedrock: Callable[[str], Awaitable[None]], +) -> None: + """Keep the opening message for a fallback to replay until Bedrock accepts the stream it opened.""" + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = message # rebind-ok: scope outlives attempts + for bedrock_message in bedrock_messages: + await send_to_bedrock(bedrock_message) + await bedrock_stream.await_output() + client_ws.scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + + +async def _ack_session_update( + client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging | None, parsed_client_message: Mapping[str, JsonValue], ) -> bool: - """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" - await bedrock_stream.await_output() - client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None) - client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + """Ack the client's session.update; False means the client is gone.""" if logging_obj is None: return True - requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( - _json_dict(parsed_client_message.get("session")).get("modalities") - ) + modalities: Final = requested_modalities(_json_dict(parsed_client_message.get("session"))) try: await client_ws.send_text( - json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities)) + json.dumps(transformation_config.session_updated_event(model, logging_obj, modalities)) ) except Exception as e: # noqa: BLE001 # any send failure means the client is gone verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) @@ -314,7 +333,7 @@ class BedrockRealtime(BaseAWSLLM): except ImportError as e: raise _sdk_import_error(self._sdk_version_lookup(), e) from e - pending_session_update: Final = _pending_session_update(websocket.scope) + pending_opening_message: Final = _pending_opening_message(websocket.scope) # Get AWS region if aws_region_name is None: @@ -376,7 +395,7 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params + websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params ) finally: await _close_bedrock_client(bedrock_client) @@ -387,7 +406,7 @@ class BedrockRealtime(BaseAWSLLM): open_bidirectional_stream: Callable[[], Awaitable[BedrockBidirectionalStream]], model: str, logging_obj: LiteLLMLogging, - pending_session_update: str | None, + pending_opening_message: str | None, drop_params: bool = False, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -396,7 +415,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") - if pending_session_update is None: + if pending_opening_message is None: await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") @@ -418,7 +437,7 @@ class BedrockRealtime(BaseAWSLLM): model, session_state, logging_obj, - initial_message=pending_session_update, + initial_message=pending_opening_message, ) logged_events: Final = ( @@ -531,18 +550,22 @@ class BedrockRealtime(BaseAWSLLM): if not await _reject_session_update(client_ws, parsed_client_message, rejection): break continue - client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, session_configuration_request=session_state.get("session_configuration_request"), ) - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message( + client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock + ) + else: + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) if is_session_update and not await _ack_session_update( - client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message + client_ws, transformation_config, model, logging_obj, parsed_client_message ): break finally: diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index e8bc0482a28..9cb38822881 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -123,26 +123,81 @@ def _integer_setting(value: JsonValue) -> int | None: return None if number is None else int(number) -def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]: +def requested_modalities(session_config: Mapping[str, JsonValue]) -> tuple[str, ...] | None: + modalities: Final = session_config.get("modalities") + if not isinstance(modalities, list): + return None + names: Final = tuple(modality for modality in modalities if isinstance(modality, str)) + return names if len(names) == len(modalities) else None + + +_OPENAI_VOICES: Final = frozenset( + ("alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse", "marin", "cedar") +) + + +def _unplayable_voice(session_config: Mapping[str, JsonValue]) -> str | None: + voice: Final = session_config.get("voice") + if not isinstance(voice, str): + return None + return voice if not voice.strip() or voice.lower() in _OPENAI_VOICES else None + + +def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]: + instructions: Final = session_config.get("instructions") + voice: Final = session_config.get("voice") + input_audio_format: Final = session_config.get("input_audio_format") + output_audio_format: Final = session_config.get("output_audio_format") + modalities: Final = session_config.get("modalities") temperature: Final = session_config.get("temperature") max_output_tokens: Final = session_config.get("max_response_output_tokens") + input_rate_hertz: Final = session_config.get("input_sample_rate_hertz") output_rate_hertz: Final = session_config.get("output_sample_rate_hertz") + tools: Final = session_config.get("tools") malformed: Final = ( - ("temperature", temperature is not None and _number_setting(temperature) is None), + ("instructions", "a string", instructions is not None and not isinstance(instructions, str)), + ("voice", "a string", voice is not None and not isinstance(voice, str)), + ("tools", "a list", tools is not None and not isinstance(tools, list)), + ( + "input_audio_format", + "a string", + input_audio_format is not None and not isinstance(input_audio_format, str), + ), + ( + "output_audio_format", + "a string", + output_audio_format is not None and not isinstance(output_audio_format, str), + ), + ( + "modalities", + "a list of strings", + modalities is not None and requested_modalities(session_config) is None, + ), + ("temperature", "a number", temperature is not None and _number_setting(temperature) is None), ( "max_response_output_tokens", + "a number or 'inf'", max_output_tokens is not None and max_output_tokens != "inf" and _integer_setting(max_output_tokens) is None, ), - ("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None), + ( + "input_sample_rate_hertz", + "a number", + input_rate_hertz is not None and _integer_setting(input_rate_hertz) is None, + ), + ( + "output_sample_rate_hertz", + "a number", + output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None, + ), ) - return tuple(name for name, is_malformed in malformed if is_malformed) + return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed) @dataclasses.dataclass(frozen=True, slots=True) class BedrockSessionUpdateRejection: - code: Literal["invalid_type", "session_update_not_applied"] + code: Literal["invalid_type", "invalid_value", "session_update_not_applied"] message: str param: str @@ -312,7 +367,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return 8000 # G.711 typically uses 8kHz return 24000 if is_output else 16000 - def transform_session_update_event(self, json_message: dict) -> list[str]: + def transform_session_update_event(self, json_message: Mapping[str, JsonValue]) -> list[str]: """ Transform session.update event to Bedrock session configuration. @@ -325,7 +380,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling session.update") messages: Final[list[str]] = [] - session_config: Final = json_message.get("session", {}) + requested_session: Final = json_message.get("session") + session_config: Final[Mapping[str, JsonValue]] = ( + requested_session if isinstance(requested_session, dict) else {} + ) locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings) # Settings Bedrock locks at prompt start may only change before it @@ -336,13 +394,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz # Update audio input configuration from session if provided - if "input_audio_format" in session_config: - input_format: Final = session_config["input_audio_format"] + input_format: Final = session_config.get("input_audio_format") + if isinstance(input_format, str): self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) # Allow direct override of the input sample rate if provided (custom extension) - if "input_sample_rate_hertz" in session_config: - self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"] + requested_input_rate: Final = _integer_setting(session_config.get("input_sample_rate_hertz")) + if requested_input_rate is not None: + self.input_sample_rate_hertz = requested_input_rate if self.prompt_started: return messages @@ -352,7 +411,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): messages.append(json.dumps(session_start)) # Send prompt start - tools: Final[Sequence[JsonValue] | None] = session_config.get("tools") + tools: Final = session_config.get("tools") prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, @@ -370,7 +429,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolUseOutputConfiguration": {"mediaType": "application/json"}, "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, } - if tools + if isinstance(tools, list) and tools else {} ), } @@ -445,12 +504,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if isinstance(requested_tools, list) else base.tools_json ) - system_prompt: Final = base.system_prompt if instructions is None else str(instructions) - voice_id: Final = base.voice_id if voice is None else str(voice) + system_prompt: Final = instructions if isinstance(instructions, str) else base.system_prompt + voice_id: Final = ( + voice if isinstance(voice, str) and _unplayable_voice(session_config) is None else base.voice_id + ) format_sample_rate: Final = ( - base.output_sample_rate_hertz - if output_audio_format is None - else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True) + self._map_audio_format_to_sample_rate(output_audio_format, is_output=True) + if isinstance(output_audio_format, str) + else base.output_sample_rate_hertz ) requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz")) output_sample_rate_hertz: Final = ( @@ -476,16 +537,26 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): session_config: Final = json_message.get("session") if not isinstance(session_config, Mapping): return None - malformed: Final = () if self.drop_params else _malformed_number_settings(session_config) + malformed: Final = () if self.drop_params else _malformed_settings(session_config) if malformed: + name, expected = malformed[0] return BedrockSessionUpdateRejection( code="invalid_type", message=( - f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature " - "and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. " + f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. " "Set drop_params to drop it instead" ), - param=f"session.{malformed[0]}", + param=f"session.{name}", + ) + unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config) + if unplayable_voice is not None: + return BedrockSessionUpdateRejection( + code="invalid_value", + message=( + f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. " + "Set drop_params to drop it instead" + ), + param="session.voice", ) if not self.prompt_started: return None @@ -785,7 +856,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, object]] = json.loads(message) + json_message: Final[dict[str, JsonValue]] = json.loads(message) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] @@ -816,11 +887,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSession: session: Final = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, - modalities=modalities if modalities is not None else ["text", "audio"], + modalities=list(modalities) if modalities is not None else ["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model @@ -842,7 +913,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSessionEvents: """Build the OpenAI session.updated ack reflecting the client's requested modalities.""" return OpenAIRealtimeStreamSessionEvents( diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 1f4e5886508..71a2dc78578 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -8,17 +8,22 @@ sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the pr Two properties are load-bearing for that billing record and are asserted by the tests: proxy identity is resolved first so it can never be evicted by caller-supplied pairs, and the whole ``user_api_key_`` prefix is reserved so a caller cannot write a proxy-authoritative -looking key. Values that break Bedrock's constraints are dropped rather than sanitised or -rejected, because an operator flipping this setting on must not turn a working request into a -400 and a silently rewritten attribution key is worse than an absent one. +looking key. Identity and ``spend_logs_metadata`` values that break Bedrock's constraints are +dropped rather than sanitised or rejected, because an operator flipping this setting on must not +turn a working request into a 400 and a silently rewritten attribution key is worse than an +absent one. For the same reason an Invoke header that would pass Bedrock's 8,500-character limit +keeps only the leading pairs that fit. A caller's own malformed ``requestMetadata`` is rejected with +a 400 unless ``drop_params`` is set, in which case its bad pairs are dropped the same way, with or +without the setting. """ from __future__ import annotations import json import re -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sized +from dataclasses import dataclass +from typing import Final, Literal import litellm @@ -26,15 +31,27 @@ BEDROCK_REQUEST_METADATA_HEADER: Final = "X-Amzn-Bedrock-Request-Metadata" BEDROCK_REQUEST_METADATA_MAX_PAIRS: Final = 16 BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX: Final = "user_api_key_" BEDROCK_REQUEST_METADATA_CLIENT_FIELD: Final = "spend_logs_metadata" +BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS: Final = 8500 _METADATA_PARAM_NAMES: Final[tuple[str, ...]] = ("metadata", "litellm_metadata") -_KEY_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$") -_VALUE_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$") +_ALLOWED_CHARACTERS: Final = r"[a-zA-Z0-9\s:_@$#=/+,.-]" +_KEY_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{1,256}}", re.ASCII) +_VALUE_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{0,256}}", re.ASCII) _OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),)) +@dataclass(frozen=True, slots=True) +class BedrockRequestMetadataViolation: + kind: Literal["type", "constraint"] + message: str + + def _is_forwardable(key: str, value: str) -> bool: - return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None + return _KEY_PATTERN.fullmatch(key) is not None and _VALUE_PATTERN.fullmatch(value) is not None + + +def _is_reserved(key: str) -> bool: + return key.strip().lower().startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) def _text_pairs(source: object) -> tuple[tuple[str, str], ...]: @@ -99,13 +116,16 @@ def _client_pairs( candidates: Final = tuple( (key, value) for key, value in (*_text_pairs(caller_metadata), *spend_logs_pairs) - if not key.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) and _is_forwardable(key, value) + if not _is_reserved(key) and _is_forwardable(key, value) ) - return tuple( - pair - for index, pair in enumerate(candidates) - if pair[0] not in tuple(earlier for earlier, _ in candidates[:index]) - )[:budget] + first_index: Final = {key: index for index, (key, _) in reversed(tuple(enumerate(candidates)))} + return tuple(pair for index, pair in enumerate(candidates) if first_index[pair[0]] == index)[:budget] + + +def _forwardable_caller_pairs(caller_metadata: object) -> tuple[tuple[str, str], ...]: + return tuple(pair for pair in _text_pairs(caller_metadata) if _is_forwardable(*pair))[ + :BEDROCK_REQUEST_METADATA_MAX_PAIRS + ] def resolve_bedrock_request_metadata( @@ -113,17 +133,18 @@ def resolve_bedrock_request_metadata( caller_metadata: object = None, ) -> dict[str, str] | None: """ - Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when the feature is - off or nothing survives Bedrock's constraints. The result is a plain dict because it is - written straight onto the Converse body, which Bedrock types as ``dict[str, str]``. + Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when nothing survives + Bedrock's constraints. The result is a plain dict because it is written straight onto the + Converse body, which Bedrock types as ``dict[str, str]``. - ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. It has already - been validated (and rejected with a 400) by the Converse transformation, so it is only - filtered here for the reserved identity prefix and the remaining slot budget. + ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. Callers reject a + malformed one with a 400 unless ``drop_params`` is set, so here its pairs are only filtered for + Bedrock's rules, the reserved identity prefix, and the remaining slot budget. With the feature + off only Bedrock's rules apply, since no identity is written for a caller to forge. """ allowed_fields: Final = _allowed_fields() if not allowed_fields: - return None + return dict(_forwardable_caller_pairs(caller_metadata)) or None sources: Final = _metadata_sources(litellm_params) identity: Final = _identity_pairs(sources, allowed_fields) client: Final = _client_pairs( @@ -136,6 +157,36 @@ def resolve_bedrock_request_metadata( return resolved or None +def _pair_violation(key: str, value: str) -> BedrockRequestMetadataViolation | None: + if _KEY_PATTERN.fullmatch(key) is None: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata key {key!r} must be 1-256 characters of {_ALLOWED_CHARACTERS}" + ) + if _VALUE_PATTERN.fullmatch(value) is None: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata value for key {key!r} must be 0-256 characters of {_ALLOWED_CHARACTERS}" + ) + return None + + +def bedrock_request_metadata_violation(caller_metadata: object) -> BedrockRequestMetadataViolation | None: + """Why a caller's ``requestMetadata`` breaks Bedrock's rules, in the order Bedrock checks them: a + wrong JSON type before the pair count and character rules. ``None`` means absent or valid.""" + text_pairs: Final = _text_pairs(caller_metadata) + entry_count: Final = len(caller_metadata) if isinstance(caller_metadata, Sized) else 0 + if caller_metadata is None: + return None + if not isinstance(caller_metadata, Mapping): + return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary") + if len(text_pairs) != entry_count: + return BedrockRequestMetadataViolation("type", "requestMetadata keys and values must be strings") + if len(text_pairs) > BEDROCK_REQUEST_METADATA_MAX_PAIRS: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata can contain a maximum of {BEDROCK_REQUEST_METADATA_MAX_PAIRS} items" + ) + return next((violation for key, value in text_pairs if (violation := _pair_violation(key, value))), None) + + def bedrock_request_metadata_is_owned() -> bool: """ Whether the proxy OWNS the request-metadata field and header name for this request. @@ -149,25 +200,51 @@ def bedrock_request_metadata_is_owned() -> bool: return bool(_allowed_fields()) +def caller_request_metadata_violations_are_dropped(drop_params: bool) -> bool: + return drop_params or litellm.drop_params is True + + +def bedrock_request_metadata_is_rewritten(drop_params: bool) -> bool: + return bedrock_request_metadata_is_owned() or caller_request_metadata_violations_are_dropped(drop_params) + + +def _header_value(pairs: Mapping[str, str]) -> str: + return json.dumps(pairs, separators=(",", ":")) + + +def _fits_header(pairs: Mapping[str, str]) -> bool: + return len(_header_value(pairs)) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + + +def _longest_fitting_prefix(resolved: Mapping[str, str]) -> dict[str, str]: + pairs: Final = tuple(resolved.items()) + return next( + (dict(pairs[:count]) for count in range(len(pairs), 0, -1) if _fits_header(dict(pairs[:count]))), + {}, + ) + + def bedrock_request_metadata_headers( litellm_params: Mapping[str, object] | None, ) -> tuple[frozenset[str], tuple[tuple[str, str], ...]]: """ - The signed ``X-Amzn-Bedrock-Request-Metadata`` header for the Invoke paths, which have no - body field for request metadata. + The signed ``X-Amzn-Bedrock-Request-Metadata`` header for an Invoke request whose body carries + no ``requestMetadata`` of its own. Returns the header names the proxy OWNS and, separately, the pairs to send. Ownership is reported whenever forwarding is enabled, including when nothing resolves, because a caller can suppress the resolver (omit the allow-listed fields, or send values that all fail Bedrock's rules) and an owned-but-empty result must still evict the caller's header rather - than fall back to it. + than fall back to it. Only the leading pairs that fit Bedrock's header limit are sent, so + identity always survives and long ``spend_logs_metadata`` pairs drop off the end. """ if not bedrock_request_metadata_is_owned(): return frozenset(), () resolved: Final = resolve_bedrock_request_metadata(litellm_params) - if resolved is None: + fitting: Final = {} if resolved is None else _longest_fitting_prefix(resolved) + if not fitting: return _OWNED_HEADER_NAMES, () - return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),) + return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, _header_value(fitting)),) def merge_bedrock_invoke_headers( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8039b85d964..82878a1744a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4,7 +4,7 @@ import json import ssl from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager -from functools import lru_cache, partial +from functools import lru_cache, partial, reduce from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +35,7 @@ import litellm.litellm_core_utils import litellm.types import litellm.types.utils from litellm._logging import redact_string, verbose_logger -from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, update_headers_with_filtered_beta from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs @@ -49,6 +49,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( synthesize_subtitle_document, ) from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params +from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import ( @@ -215,6 +216,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -1515,8 +1517,65 @@ class BaseLLMHTTPHandler: api_base: str, api_key: str | None, headers: Mapping[str, str], - ) -> tuple[str, dict[str, str], dict[str, object]]: - outbound_headers: Final = provider_config.validate_environment(headers=headers, model=model, api_key=api_key) + litellm_params: Mapping[str, object], + ) -> tuple[str, Mapping[str, str], Mapping[str, object]]: + outbound_headers: Final = provider_config.validate_environment( + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ) + return self._prepare_decisions_request_with_headers( + model=model, + logging_obj=logging_obj, + provider_config=provider_config, + body=body, + logged_body=logged_body, + api_base=api_base, + api_key=api_key, + outbound_headers=outbound_headers, + ) + + async def _aprepare_decisions_request( + self, + model: str, + logging_obj: LiteLLMLoggingObj | None, + provider_config: BaseDecisionsConfig, + body: Mapping[str, object], + logged_body: Mapping[str, object], + api_base: str, + api_key: str | None, + headers: Mapping[str, str], + litellm_params: Mapping[str, object], + ) -> tuple[str, Mapping[str, str], Mapping[str, object]]: + outbound_headers: Final = await provider_config.avalidate_environment( + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ) + return self._prepare_decisions_request_with_headers( + model=model, + logging_obj=logging_obj, + provider_config=provider_config, + body=body, + logged_body=logged_body, + api_base=api_base, + api_key=api_key, + outbound_headers=outbound_headers, + ) + + def _prepare_decisions_request_with_headers( + self, + model: str, + logging_obj: LiteLLMLoggingObj | None, + provider_config: BaseDecisionsConfig, + body: Mapping[str, object], + logged_body: Mapping[str, object], + api_base: str, + api_key: str | None, + outbound_headers: Mapping[str, str], + ) -> tuple[str, Mapping[str, str], Mapping[str, object]]: url: Final = provider_config.get_complete_url(api_base=api_base, model=model) data: Final = dict(body) if logging_obj is not None: @@ -1553,6 +1612,7 @@ class BaseLLMHTTPHandler: api_base=api_base, api_key=api_key, headers=headers, + litellm_params=litellm_params, ) signed_headers, signed_body = provider_config.sign_request( headers=outbound_headers, @@ -1567,7 +1627,7 @@ class BaseLLMHTTPHandler: response: Final = sync_httpx_client.post( url, content=signed_body, - json=data if signed_body is None else None, + json=dict(data) if signed_body is None else None, headers=dict(signed_headers), timeout=timeout, logging_obj=logging_obj, @@ -1594,7 +1654,7 @@ class BaseLLMHTTPHandler: litellm_params: Mapping[str, object], client: AsyncHTTPHandler | None = None, ) -> DecisionsIRResponse: - url, outbound_headers, data = self._prepare_decisions_request( + url, outbound_headers, data = await self._aprepare_decisions_request( model=model, logging_obj=logging_obj, provider_config=provider_config, @@ -1603,6 +1663,7 @@ class BaseLLMHTTPHandler: api_base=api_base, api_key=api_key, headers=headers, + litellm_params=litellm_params, ) sign: Final = partial( provider_config.sign_request, @@ -1625,7 +1686,7 @@ class BaseLLMHTTPHandler: response: Final = await async_httpx_client.post( url, content=signed_body, - json=data if signed_body is None else None, + json=dict(data) if signed_body is None else None, headers=dict(signed_headers), timeout=timeout, logging_obj=logging_obj, @@ -2245,6 +2306,16 @@ class BaseLLMHTTPHandler: merged_headers.update(extra_headers_from_kwargs) if provider_specific_headers: merged_headers.update(provider_specific_headers) + request_params: Final = reduce( + delete_nested_value, + litellm_params.get("additional_drop_params") or (), + anthropic_messages_optional_request_params, + ) + if ( + request_params.get("safeguards") is not None + and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() + ): + await aload_beta_headers_config() ( headers, api_base, @@ -2252,7 +2323,7 @@ class BaseLLMHTTPHandler: headers=merged_headers or {}, model=model, messages=messages, - optional_params=anthropic_messages_optional_request_params, + optional_params=request_params, litellm_params=dict(litellm_params), api_key=api_key, api_base=api_base, @@ -2283,20 +2354,11 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, ) - additional_drop_params: Final[list[str]] = litellm_params.get("additional_drop_params") or [] - if additional_drop_params: - from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value - - for path in additional_drop_params: - anthropic_messages_optional_request_params = delete_nested_value( - anthropic_messages_optional_request_params, path - ) - # Prepare request body request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request( model=model, messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -2403,7 +2465,7 @@ class BaseLLMHTTPHandler: held_back_tool_names: Final = self._server_fulfilled_tools_in_request( logging_obj=logging_obj, - tools=anthropic_messages_optional_request_params.get("tools"), + tools=request_params.get("tools"), ) initial_response = AgenticAnthropicStreamingIterator( completion_stream=completion_stream, @@ -2411,7 +2473,7 @@ class BaseLLMHTTPHandler: model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs=kwargs_for_agentic, @@ -2434,7 +2496,7 @@ class BaseLLMHTTPHandler: model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs=kwargs_for_agentic, @@ -6807,6 +6869,7 @@ class BaseLLMHTTPHandler: first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, turn_admission: Callable[[], Awaitable[None]] | None = None, + frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: """ @@ -6822,6 +6885,35 @@ class BaseLLMHTTPHandler: """ _ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks() + _ws_guardrail_callbacks: list = [] + _ws_output_guardrail_callbacks: list = [] + try: + import litellm as _litellm + + # Use duck-typing so any guardrail that exposes the PII + # masking interface works, not just _OPTIONAL_PresidioPIIMasking. + # This avoids a layering violation (SDK importing from proxy). + _ws_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) + and callable(getattr(cb, "_unmask_pii_text", None)) + and getattr(cb, "output_parse_pii", False) + ] + _ws_output_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) + and getattr(cb, "apply_to_output", False) + ] + except Exception as _guardrail_exc: + verbose_logger.warning( + "Responses WebSocket: failed to collect guardrail callbacks — PII masking will be skipped. Error: %s", + _guardrail_exc, + ) + if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, @@ -6840,6 +6932,8 @@ class BaseLLMHTTPHandler: first_message=first_message, quota_callbacks=_ws_quota_callbacks, turn_admission=turn_admission, + guardrail_callbacks=_ws_guardrail_callbacks, # pyright: ignore[reportUnknownArgumentType] # duck-typed callbacks from litellm.callbacks + frame_guard=frame_guard, **kwargs, ) return await handler.run() @@ -6920,36 +7014,6 @@ class BaseLLMHTTPHandler: if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata - _ws_guardrail_callbacks: list = [] - _ws_output_guardrail_callbacks: list = [] - try: - import litellm as _litellm - - # Use duck-typing so any guardrail that exposes the PII - # masking interface works, not just _OPTIONAL_PresidioPIIMasking. - # This avoids a layering violation (SDK importing from proxy). - _ws_guardrail_callbacks = [ - cb - for cb in _litellm.callbacks - if callable(getattr(cb, "check_pii", None)) - and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) - and callable(getattr(cb, "_unmask_pii_text", None)) - and getattr(cb, "output_parse_pii", False) - ] - _ws_output_guardrail_callbacks = [ - cb - for cb in _litellm.callbacks - if callable(getattr(cb, "check_pii", None)) - and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) - and getattr(cb, "apply_to_output", False) - ] - except Exception as _guardrail_exc: - verbose_logger.warning( - "Responses WebSocket: failed to collect guardrail " - "callbacks — PII masking will be skipped. Error: %s", - _guardrail_exc, - ) - streaming: Final = ResponsesWebSocketStreaming( websocket=websocket, backend_ws=cast(ClientConnection, backend_ws), @@ -6964,6 +7028,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, turn_admission=turn_admission, + frame_guard=frame_guard, ) return await streaming.bidirectional_forward() diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index aa3d7942174..2acf1c69594 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -87,7 +87,7 @@ def _keyed_rows(model: str, quality: str) -> tuple[tuple[int, int, float], ...]: suffix: Final = f"/{model}" return tuple( (width, height, float(raw_cost)) - for key in litellm.model_cost + for key in (*litellm.model_cost,) if isinstance(key, str) and key.startswith(prefix) and key.endswith(suffix) for size in (key[len(prefix) : -len(suffix)],) for dimensions in (_parse_keyed_dimensions(size),) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 8b6946ee562..14ae59d4f85 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -544,7 +544,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return cached[2] index: Final[list[tuple[str, dict]]] = [] - for key, model_info in model_cost.items(): + for key, model_info in (*model_cost.items(),): if not key.startswith("fireworks_ai/"): continue if not isinstance(model_info, dict): diff --git a/litellm/llms/nace/decisions/transformation.py b/litellm/llms/nace/decisions/transformation.py new file mode 100644 index 00000000000..c96177bb7d5 --- /dev/null +++ b/litellm/llms/nace/decisions/transformation.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class NaceDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("NACE_API_KEY",) + api_base_env = ("NACE_API_BASE", "NACE_BASE_URL") + + def get_default_api_base(self) -> str | None: + return "https://console.nace.ai" diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index b5328b1ba4d..284a6ffea62 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -197,7 +197,7 @@ class OllamaModelInfo(BaseLLMModelInfo): "ollama/" + stripped_model, "ollama_chat/" + stripped_model, } - model_cost_keys: Final = {key.lower() for key in model_cost} + model_cost_keys: Final = {key.lower() for key in (*model_cost,)} return any(name.lower() in model_cost_keys for name in potential_model_names) @staticmethod diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 94c5c5d016d..5bf19844287 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -28,7 +28,9 @@ def _catalogue_declares_default_effort() -> bool: the one litellm gave before it existed. Scanning costs ~80us on the largest published map and only on the fallback path, which is noise beside the request it precedes. """ - return any(isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in litellm.model_cost.values()) + return any( + isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in (*litellm.model_cost.values(),) + ) def _normalize_reasoning_effort_for_chat_completion( diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 0c1eaaaae26..73f40f5e0e4 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -11,13 +11,14 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer import openai -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream from openai._base_client import make_request_options from openai._constants import RAW_RESPONSE_HEADER from openai._legacy_response import LegacyAPIResponse from openai._types import RequestOptions from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted +from openai.types.chat import ChatCompletion, ChatCompletionChunk from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, TypeAdapter from typing_extensions import overload @@ -354,7 +355,14 @@ _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body")) -def _embedding_request_without_sdk_defaults( +def _raw_chat_response( + result: ChatCompletion | Stream[ChatCompletionChunk] | AsyncStream[ChatCompletionChunk], +) -> LegacyAPIResponse[ChatCompletion]: + assert isinstance(result, LegacyAPIResponse) + return result + + +def _request_without_sdk_transform( data: Mapping[str, object], timeout: float | httpx.Timeout ) -> tuple[Mapping[str, object], RequestOptions]: body: Final = {k: v for k, v in data.items() if k not in _SDK_OPTION_KEYS} @@ -517,7 +525,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ start_time: Final = time.time() try: - raw_response = await openai_aclient.chat.completions.with_raw_response.create(**data, timeout=timeout) + body, options = _request_without_sdk_transform(data, timeout) + raw_response = _raw_chat_response( + await openai_aclient.post( + "/chat/completions", + body=body, + cast_to=ChatCompletion, + options=options, + stream=bool(data.get("stream")), + stream_cls=AsyncStream[ChatCompletionChunk], + ) + ) end_time = time.time() if hasattr(raw_response, "headers"): @@ -558,7 +576,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ raw_response = None try: - raw_response = openai_client.chat.completions.with_raw_response.create(**data, timeout=timeout) + body, options = _request_without_sdk_transform(data, timeout) + raw_response = _raw_chat_response( + openai_client.post( + "/chat/completions", + body=body, + cast_to=ChatCompletion, + options=options, + stream=bool(data.get("stream")), + stream_cls=Stream[ChatCompletionChunk], + ) + ) if hasattr(raw_response, "headers"): headers = dict(raw_response.headers) @@ -1257,7 +1285,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, ) -> LegacyAPIResponse[CreateEmbeddingResponse]: if "encoding_format" not in data: - body, options = _embedding_request_without_sdk_defaults(data, timeout) + body, options = _request_without_sdk_transform(data, timeout) bypass_response: Final = await openai_aclient.post( "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse ) @@ -1274,7 +1302,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, ) -> LegacyAPIResponse[CreateEmbeddingResponse]: if "encoding_format" not in data: - body, options = _embedding_request_without_sdk_defaults(data, timeout) + body, options = _request_without_sdk_transform(data, timeout) bypass_response: Final = openai_client.post( "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse ) diff --git a/litellm/llms/openai_like/chat/transformation.py b/litellm/llms/openai_like/chat/transformation.py index 83fcb43c9b8..e07b5627256 100644 --- a/litellm/llms/openai_like/chat/transformation.py +++ b/litellm/llms/openai_like/chat/transformation.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage from litellm.types.utils import ModelResponse @@ -119,6 +120,7 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): choice["message"] = message returned_response: Final = ModelResponse(**response_json) + set_provider_response_headers_in_hidden_params(returned_response, response.headers) if custom_llm_provider is not None: returned_response.model = custom_llm_provider + "/" + (returned_response.model or "") diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 2e9a300e2fd..94f9792da35 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -8,6 +8,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -57,6 +58,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): headers=normalized, optional_params=optional_params, messages=messages, + beta_map_provider=LlmProviders.OPENAI_LIKE.value, ) return merged, api_base diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index b0da53af8b3..885fc98ffea 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -6,17 +6,21 @@ It uses the field targeting configuration from litellm_logging_obj to extract specific fields for guardrail processing. """ -from collections.abc import Mapping +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import without_server_streaming_classification -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, + StreamTransformSink, +) from litellm.types.passthrough_endpoints.pass_through_endpoints import PassThroughGuardrailSettings from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -208,7 +212,7 @@ class _EventStreamMediaTypeProvider(Protocol): @runtime_checkable class _EventStreamEndpointChecker(Protocol): - def event_stream_endpoint_is_de_anonymizable(self, endpoint: str) -> bool: ... + def event_stream_endpoint_is_de_anonymizable(self, endpoint: str, model: str, body: object) -> bool: ... _PROVIDER_HANDLERS: dict[str, type[BaseTranslation]] = {} @@ -225,6 +229,15 @@ def _get_provider_handlers() -> dict[str, type[BaseTranslation]]: return _PROVIDER_HANDLERS +def _provider_handler(provider: object) -> type[BaseTranslation] | None: + return _get_provider_handlers().get(provider) if isinstance(provider, str) else None + + +def _owns_event_stream_item(handler_cls: type[BaseTranslation], item: object) -> bool: + owns: Final[Callable[[object], bool] | None] = getattr(handler_cls, "owns_event_stream_item", None) + return owns is not None and owns(item) + + class LlmPassthroughRouteHandler(BaseTranslation): """ Dispatcher for allm_passthrough_route guardrail translation. @@ -277,6 +290,100 @@ class LlmPassthroughRouteHandler(BaseTranslation): request_data=request_data, ) + @staticmethod + def _event_stream_translation(items: Sequence[object], provider: object = None) -> BaseTranslation | None: + """The provider handler for a relayed event stream: the one that owns its frames, else ``provider``'s.""" + handlers: Final = _get_provider_handlers() + owner: Final = next( + ( + handler_cls + for handler_cls in handlers.values() + if items and _owns_event_stream_item(handler_cls, items[0]) + ), + handlers.get(provider) if isinstance(provider, str) else None, + ) + return owner() if owner is not None else None + + async def process_output_streaming_response( + self, + responses_so_far: Sequence[object], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: dict[str, object] | None = None, # mutable-ok: the request dict guardrails receive + stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, + ) -> Sequence[object]: + translation: Final = self._event_stream_translation( + responses_so_far, (request_data or {}).get("custom_llm_provider") + ) + if translation is not None: + await translation.process_output_streaming_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + responses_so_far=list(responses_so_far), + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + stream_transform_sink=stream_transform_sink, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, + ) + return responses_so_far + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + translation: Final = self._event_stream_translation(responses_so_far) + return translation.get_streaming_scan_key(responses_so_far) if translation is not None else None + + def needs_ended_stream_scan(self, guardrail: "CustomGuardrail", responses_so_far: Sequence[object]) -> bool: + translation: Final = self._event_stream_translation(responses_so_far) + return translation is not None and translation.needs_ended_stream_scan(guardrail, responses_so_far) + + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: + has_ended: Final[Callable[[Sequence[object]], bool] | None] = getattr( + self._event_stream_translation(responses_so_far), "_check_streaming_has_ended", None + ) + return has_ended is not None and has_ended(responses_so_far) + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes] | None: + translation: Final = self._event_stream_translation( + responses_so_far or (), exc.request_data.get("custom_llm_provider") + ) + if translation is None: + return None + return translation.build_block_sse_chunks(exc, stream_started=stream_started, responses_so_far=responses_so_far) + + @staticmethod + def event_stream_items( + provider: object, endpoint: object, chunks: AsyncIterable[bytes] + ) -> AsyncIterator[object] | None: + """The upstream event stream of ``endpoint`` cut into the items the provider's guardrail translation scans, or + None when the provider has no streaming guardrail translation.""" + framer: Final[Callable[[AsyncIterable[bytes], str], AsyncIterator[object]] | None] = getattr( + _provider_handler(provider), "event_stream_items", None + ) + return framer(chunks, endpoint if isinstance(endpoint, str) else "") if framer is not None else None + + @staticmethod + def encode_event_stream_item(provider: object, item: object) -> bytes: + encoder: Final[Callable[[object], bytes] | None] = getattr( + _provider_handler(provider), "encode_event_stream_item", None + ) + if encoder is None: + raise TypeError(f"provider {provider!r} has no event-stream encoder") + return encoder(item) + + @staticmethod + def event_stream_error_item(provider: object, message: str, status_code: int) -> bytes | None: + """An in-stream error frame in the provider's event-stream format, or None when the format has none.""" + builder: Final[Callable[[str, int], bytes] | None] = getattr( + _provider_handler(provider), "event_stream_error_item", None + ) + return builder(message, status_code) if builder is not None else None + @staticmethod def is_event_stream_response(provider: str | None, content_type: str) -> bool: handler_cls: Final = _get_provider_handlers().get(provider or "") @@ -297,11 +404,26 @@ class LlmPassthroughRouteHandler(BaseTranslation): return getattr(handler_cls, "de_anonymize_event_stream", None) @staticmethod - def supports_event_stream_de_anonymization(provider: str | None, endpoint: str | None) -> bool: - handler_cls: Final = _get_provider_handlers().get(provider or "") + def supports_event_stream_relay(provider: object, endpoint: object, model: object, body: object) -> bool: + """Whether the provider's guardrail translation can scan this endpoint's event stream for ``model`` and the + request ``body`` frame by frame as it is relayed.""" + relayable: Final[Callable[[str, str, object], bool] | None] = getattr( + _provider_handler(provider), "event_stream_endpoint_is_relayable", None + ) + return relayable is not None and relayable( + endpoint if isinstance(endpoint, str) else "", model if isinstance(model, str) else "", body + ) + + @staticmethod + def supports_event_stream_de_anonymization(provider: object, endpoint: object, model: object, body: object) -> bool: + """Whether the provider's guardrail translation can scan this endpoint's whole event stream for ``model`` and + the request ``body`` and rewrite its frames.""" + handler_cls: Final = _provider_handler(provider) if not isinstance(handler_cls, _EventStreamEndpointChecker): return False - return handler_cls.event_stream_endpoint_is_de_anonymizable(endpoint or "") + return handler_cls.event_stream_endpoint_is_de_anonymizable( + endpoint if isinstance(endpoint, str) else "", model if isinstance(model, str) else "", body + ) @staticmethod async def de_anonymize_event_stream( diff --git a/litellm/main.py b/litellm/main.py index 23cd3c40250..3eae64d828e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4132,6 +4132,7 @@ def _complete_sagemaker(ctx: CompletionDispatchContext) -> _CompletionDispatchRe _ADDITIONAL_DROP_PARAMS_ADAPTER: Final = TypeAdapter(list[str]) _OPTIONAL_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) _BASE_URL_HINT: Final = TypeAdapter(BaseUrlRoutingHint) @@ -5695,7 +5696,9 @@ def completion( and provider_config.accepts_reasoning_summary_in_reasoning_effort, ) optional_params = get_optional_params( - **{**optional_param_args, "reasoning_effort": provider_reasoning_effort}, **non_default_params + **{**optional_param_args, "reasoning_effort": provider_reasoning_effort}, + model_info=_MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else None, + **non_default_params, ) processed_non_default_params: Final = pre_process_non_default_params( model=model, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c3efaf168ab..6ebb54c0035 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1546,7 +1546,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -1623,7 +1623,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -1700,7 +1700,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -1777,7 +1777,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -12877,6 +12877,7 @@ "supports_prompt_caching": true }, "azure_ai/ministral-3b": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-08, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12940,6 +12941,7 @@ "supports_vision": true }, "azure_ai/mistral-medium-2505": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12977,6 +12979,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -29699,6 +29702,7 @@ }, "vertex_ai/gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -29806,6 +29810,7 @@ "supports_native_streaming": true }, "vertex_ai/gemini-3.8-live": { + "deprecation_date": "2027-09-25", "input_cost_per_audio_token": 3e-06, "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, @@ -31960,6 +31965,7 @@ }, "gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -32409,11 +32415,16 @@ ] }, "github_copilot/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32427,11 +32438,16 @@ "thinking_always_on": true }, "github_copilot/claude-fable-5.1": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32490,11 +32506,16 @@ "thinking_always_on": false }, "github_copilot/claude-opus-4.8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32524,13 +32545,17 @@ "supports_vision": true }, "github_copilot/claude-opus-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, "deprecation_date": "2026-11-16", + "input_cost_per_token": 5e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", + "output_cost_per_token": 2.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32543,11 +32568,16 @@ "supports_vision": true }, "github_copilot/claude-opus-5.5": { + "cache_creation_input_token_cost": 5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 2e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32561,11 +32591,16 @@ "thinking_always_on": true }, "github_copilot/claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32578,11 +32613,16 @@ "supports_vision": true }, "github_copilot/claude-sonnet-5.5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32596,11 +32636,15 @@ "thinking_always_on": true }, "github_copilot/gemini-3.7-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions" ], @@ -32610,11 +32654,15 @@ "supports_vision": true }, "github_copilot/gemini-3.8-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions" ], @@ -32780,11 +32828,18 @@ "supports_vision": true }, "github_copilot/gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -32796,11 +32851,15 @@ "supports_vision": true }, "github_copilot/gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 4.5e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32811,11 +32870,18 @@ "supports_vision": true }, "github_copilot/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32826,11 +32892,20 @@ "supports_vision": true }, "github_copilot/gpt-5.6-luna": { + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_200k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_200k_tokens": 4e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_200k_tokens": 4e-07, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32841,11 +32916,20 @@ "supports_vision": true }, "github_copilot/gpt-5.6-sol": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32856,13 +32940,21 @@ "supports_vision": true }, "github_copilot/gpt-5.6-terra": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, "deprecation_date": "2026-11-16", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32873,11 +32965,20 @@ "supports_vision": true }, "github_copilot/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32888,11 +32989,20 @@ "supports_vision": true }, "github_copilot/gpt-6-luna": { + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "input_cost_per_token": 1e-07, + "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 5e-07, + "output_cost_per_token_above_272k_tokens": 7.5e-07, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32903,11 +33013,20 @@ "supports_vision": true }, "github_copilot/gpt-6-sol": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32918,11 +33037,20 @@ "supports_vision": true }, "github_copilot/gpt-6.1-sol": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_above_272k_tokens": 2e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32933,11 +33061,15 @@ "supports_vision": true }, "github_copilot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, "litellm_provider": "github_copilot", "max_input_tokens": 917504, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions" ], @@ -32947,11 +33079,15 @@ "supports_vision": true }, "github_copilot/mai-code-1.1-flash": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, "litellm_provider": "github_copilot", "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1.2e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -45295,6 +45431,38 @@ "output_cost_per_token": 0.0, "source": "https://docs.perplexity.ai/docs/getting-started/pricing" }, + "perplexity/pplx-decider-v1.2": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-latest": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-v1.2": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-latest": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, "perplexity/sonar": { "input_cost_per_token": 1e-06, "litellm_provider": "perplexity", @@ -65431,6 +65599,40 @@ "supports_vision": false, "source": "https://deepinfra.com/pricing" }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash" + }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash-1017" + }, "deepinfra/openai/gpt-oss-120b-Ultra": { "max_tokens": 131072, "max_input_tokens": 131072, @@ -71184,6 +71386,38 @@ "notes": "Self-hosted decision model; infrastructure costs are paid separately" } }, + "nace/drex-v1.0": { + "input_cost_per_token": 4e-08, + "litellm_provider": "nace", + "max_input_tokens": 32768, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v1.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v2.0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-latest": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", @@ -72466,7 +72700,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72499,7 +72733,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72856,6 +73090,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72887,6 +73122,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72933,6 +73169,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73003,6 +73240,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -73039,6 +73277,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -73082,6 +73321,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73123,6 +73363,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73164,6 +73405,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73286,6 +73528,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -73329,6 +73572,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74301,6 +74545,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74337,6 +74582,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74388,6 +74634,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74468,6 +74715,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -74509,6 +74757,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -74557,6 +74806,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74603,6 +74853,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74649,6 +74900,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74699,6 +74951,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -74747,6 +75000,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -75141,6 +75395,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-4-argon": { + "deprecation_date": "2027-10-13", "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 420f42c690c..744f9e1aa4c 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -99,7 +99,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): self._litellm_logging_obj = litellm_logging_obj self._spend = _SpendCollection(provider_config, litellm_logging_obj) self._flush_scheduled = False - self._background_tasks: set[asyncio.Task] = set() # mutable-ok: instance set for background task tracking + self._background_tasks: set[asyncio.Task[bool]] = set() # mutable-ok: instance set for background task tracking + self._flush: asyncio.Task[bool] | None = None self._hidden_params: dict[str, object] = {} # mutable-ok: router attaches response headers here in place @property @@ -157,6 +158,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): ) self._background_tasks.add(task) + self._flush = task task.add_done_callback(self._background_tasks.discard) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging @@ -166,6 +168,16 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): e, ) + async def wait_for_spend_flush(self) -> bool: + """Wait until the spend-tracking flush scheduled when the stream ended or was closed has finished, and say + whether it logged the stream, which it does not when the stream returned nothing it could read. Cancelling the + wait leaves the flush running.""" + flush: Final = self._flush + if flush is None: + return False + outcome: Final = await asyncio.shield(asyncio.gather(flush, return_exceptions=True)) + return outcome[0] is True + def __aiter__(self) -> AsyncPassthroughStreamingResponse: return self diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 449adaa73e9..3b2c2e09df5 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -862,12 +862,17 @@ def _raise_unless_oauth2_discovery_server( def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool: """True when a DCR-bridge server relays client registration to the upstream authorization - server instead of short-circuiting to an admin-configured OAuth client. In the relay arm the + server instead of short-circuiting to the server's own OAuth client. In the relay arm the upstream holds each client's own registration, so the authorize and token relays pass the client's ``client_id`` and ``redirect_uri`` through verbatim and the authorization code returns directly to the client's redirect URI without transiting the gateway. Gateway-side redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit - arm, where the upstream only knows the gateway's own callback.""" + arm, where the upstream only knows the gateway's own callback. + + Authorize and token decide per request from the server's current client: once the gateway + holds its shared client (registered at ``/register`` under ``PROXY_BASE_URL``, or entered by an + admin) every caller runs the gateway-callback arm, while a server that still has none keeps + relaying, so clients registered before the shared client existed keep their own binding.""" return bool( mcp_server.is_dcr_bridge and mcp_server.effective_registration_url @@ -880,6 +885,27 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool: ) +def _register_relays_bridge_client(mcp_server: MCPServer) -> bool: + """``/register`` relays the caller's own registration only when the gateway cannot keep a shared + client: with ``PROXY_BASE_URL`` set it registers one client bound to its own ``/callback`` and + keeps it as the server's client, so an upstream that restricts redirect domains (Atlassian) + never sees a laptop loopback redirect. Without a database the shared client would live only in + this process, so a database-free gateway keeps relaying as well.""" + return _dcr_bridge_relays_client_registration(mcp_server) and ( + configured_proxy_base_url() is None or not _shared_client_store_available() + ) + + +def _shared_client_store_available() -> bool: + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import + + try: + get_prisma_client_or_throw("Database not connected. Cannot keep the shared MCP OAuth client.") + except HTTPException: + return False + return True + + def _require_s256_pkce( code_challenge: str | None, code_challenge_method: str | None, @@ -1028,6 +1054,17 @@ async def _managed_authorization_scope( return selected_scope, scope_user_id +async def _hydrate_bridge_shared_client(mcp_server: MCPServer, request: Request) -> None: + """Whichever worker served ``/register`` persisted the bridge's shared gateway client; a worker + whose catalog snapshot predates that write still holds no ``client_id``, so the gateway-callback + arms load the persisted client before presenting one upstream. Callers that carry their own + client (a gateway-minted ephemeral client sealed into the state or the code) skip this, since the + upstream bound their code to that client.""" + if not mcp_server.is_dcr_bridge or mcp_server.client_id: + return + await _reuse_persisted_dcr_client_if_available(mcp_server, f"{get_request_base_url(request)}/callback") + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -1069,6 +1106,8 @@ async def authorize_with_server( # calling this for its enforcement side effect, then falls through to the gateway # /callback flow below, which reads the original code_challenge names. bridge_challenge, bridge_method = _require_s256_pkce(code_challenge, code_challenge_method) + if ephemeral_dcr_client is None: + await _hydrate_bridge_shared_client(resolved_server, request) # A gateway-minted ephemeral client is registered against {base}/callback, so its # flow must run the short-circuit arm; the relay arm is only for clients that # registered themselves through the front door and hold their own redirect binding. @@ -1233,6 +1272,7 @@ async def exchange_token_with_server( refresh_token: str | None = None, scope: str | None = None, client_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, + sealed_client: bool = False, ): raise_if_not_oauth2(mcp_server) if grant_type not in ("authorization_code", "refresh_token"): @@ -1282,11 +1322,14 @@ async def exchange_token_with_server( # recovered from a sealed code) must authenticate the way its own registration was granted, # not the way the server row is configured; callers that carry no method keep the row's method # as before. - resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id - resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret + if not sealed_client: + await _hydrate_bridge_shared_client(resolved_server, request) + server_client_wins: Final = bool(resolved_server.client_id) and not sealed_client + resolved_client_id: Final = resolved_server.client_id if server_client_wins else client_id + resolved_client_secret: Final = resolved_server.client_secret if server_client_wins else client_secret resolved_auth_method: Final = ( resolved_server.token_endpoint_auth_method - if resolved_server.client_id + if server_client_wins else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method) ) try: @@ -1830,10 +1873,13 @@ async def _persist_dcr_client_registration( empty and otherwise adopts whatever a concurrent registration persisted first. The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped - unconditionally: the caller holds the upstream token and the gateway must hold no OAuth - client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a - ``client_id`` onto a server whose mode promises the gateway stores nothing, making a - fresh pass-through server read as gateway-authorized. + unless the server is a DCR bridge: the caller holds the upstream token and the gateway holds + no OAuth client identity for these servers, so persisting would stamp ``oauth2_flow`` and a + ``client_id`` onto a server whose mode promises the gateway stores nothing. A DCR bridge is + the exception by design: its gateway-hosted sign-in registers one shared client bound to the + gateway ``/callback`` (the same app an admin may enter by hand), and that client must survive + for every later caller's authorize and token exchange. ``oauth2_flow`` is still never stamped + on a client-forwarded row. ``redirect_uris`` records what the client is bound to so a later origin change can be detected as a positive mismatch and trigger re-registration instead of stranding the @@ -1842,7 +1888,7 @@ async def _persist_dcr_client_registration( ``update_mcp_server`` merges credential blobs: a re-registered public client must not inherit the previous client's secret or auth method. """ - if mcp_server.is_client_forwarded_token: + if mcp_server.is_client_forwarded_token and not mcp_server.is_dcr_bridge: return "skipped" try: @@ -1910,20 +1956,23 @@ async def _persist_dcr_client_registration( ): await _refresh_persisted_dcr_server(stored) return "reused" + stamped_flow: Final[Literal["authorization_code"] | None] = ( + None if mcp_server.is_client_forwarded_token else "authorization_code" + ) updated_row: Final = await update_mcp_server( prisma_client=prisma_client, data=( UpdateMCPServerRequest( server_id=mcp_server.server_id, credentials=credentials, - oauth2_flow="authorization_code", + oauth2_flow=stamped_flow, token_url=mcp_server.token_url, ) if mcp_server.token_url else UpdateMCPServerRequest( server_id=mcp_server.server_id, credentials=credentials, - oauth2_flow="authorization_code", + oauth2_flow=stamped_flow, ) ), touched_by="mcp_oauth_dcr", @@ -2211,7 +2260,7 @@ async def register_client_with_server( ) return dummy_return - bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server) + bridge_relay: Final = _register_relays_bridge_client(resolved_server) if bridge_relay and not client_redirect_uris: raise HTTPException( status_code=400, diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 90289fd4c20..70e4bac01b5 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -248,6 +248,7 @@ class SubjectTokenRefusal(LiteLLMBaseModel): "invalid_grant", "invalid_client", "invalid_scope", + "invalid_target", "temporarily_unavailable", ] description: str = Field(min_length=1) @@ -1527,7 +1528,7 @@ def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response: return _oauth_error(503, refusal.error, refusal.description) case "invalid_client": return _oauth_error(401, refusal.error, refusal.description) - case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope": + case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope" | "invalid_target": return _oauth_error(400, refusal.error, refusal.description) case _: assert_never(refusal.error) diff --git a/litellm/proxy/_experimental/mcp_server/id_jag_grant.py b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py index 155d5280f5a..f2177df5a26 100644 --- a/litellm/proxy/_experimental/mcp_server/id_jag_grant.py +++ b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py @@ -23,6 +23,7 @@ from urllib.parse import unquote_plus import jwt from fastapi import Request +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.models.user import LiteLLM_UserTable @@ -48,6 +49,7 @@ _INACTIVE_SUBJECT: Final = frozenset(("no_active_key", "unresolvable")) _REPLAY_KEY_PREFIX: Final = "mcp_gateway_dcr:id_jag:" SCOPE_MCP: Final = "litellm:mcp" SCOPE_LLM: Final = "litellm:llm" +_RESOURCE_CLAIM_ADAPTER: Final = TypeAdapter(list[str]) SCOPE_ROUTE_GROUPS: Final = MappingProxyType( {SCOPE_MCP: ("mcp_inference_routes",), SCOPE_LLM: ("llm_api_routes_without_mcp",)} ) @@ -65,9 +67,9 @@ class GrantedScope: def granted_scope(claim: object, *, mcp_resource: bool) -> GrantedScope | SubjectTokenRefusal: """The routes the assertion buys. With ``aud`` the bare origin the ``scope`` claim decides and - no claim is the full credential. With ``aud`` the MCP resource (``{base}/mcp``) the identity - provider named that resource, so the credential is ceilinged to the MCP routes whatever the - claim says: ``litellm:llm`` cannot be granted and is refused.""" + no claim is the full credential. With the MCP resource (``{base}/mcp``) as ``aud`` or as every + ``resource``, the identity provider named that resource, so the credential is ceilinged to the + MCP routes whatever the claim says: ``litellm:llm`` cannot be granted and is refused.""" if claim is not None and not isinstance(claim, str): return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION) requested: Final = frozenset(claim.split()) if claim else frozenset() @@ -90,6 +92,32 @@ def names_mcp_resource(aud: object, base_url: str | None) -> bool: return base_url is not None and f"{base_url}/mcp" in audiences and base_url not in audiences +def _resource_claim_values(claim: object) -> tuple[str, ...] | None: + if isinstance(claim, str): + return (claim,) + if not isinstance(claim, list): + return None + try: + values: Final = _RESOURCE_CLAIM_ADAPTER.validate_python(claim, strict=True) + except ValidationError: + return None + return tuple(values) if values else None + + +def resource_names_mcp(claims: Mapping[str, object], base_url: str | None) -> bool | SubjectTokenRefusal: + """The signed RFC 8707 ``resource`` claim (a URI or an array of URIs) restricts the credential like an + ``aud`` of ``{base}/mcp`` does.""" + if "resource" not in claims: + return False + resources: Final = _resource_claim_values(claims["resource"]) + if resources is None: + return SubjectTokenRefusal(error="invalid_target", description=_RESOURCE_DESCRIPTION) + served_resources: Final = frozenset(gateway_audiences(base_url)) + if any(resource not in served_resources for resource in resources): + return SubjectTokenRefusal(error="invalid_target", description=_RESOURCE_DESCRIPTION) + return frozenset(resources) == frozenset((f"{base_url}/mcp",)) + + def _route_groups(scopes: Sequence[str]) -> Iterator[str]: for scope in scopes: yield from SCOPE_ROUTE_GROUPS[scope] @@ -98,8 +126,10 @@ def _route_groups(scopes: Sequence[str]) -> Iterator[str]: _SCOPE_DESCRIPTION: Final = ( f"scope must be {SCOPE_MCP} (MCP gateway routes), {SCOPE_LLM} (LLM API routes), or both, space separated" ) +_RESOURCE_DESCRIPTION: Final = "the assertion's resource claim must name this gateway's origin or its /mcp resource" _MCP_RESOURCE_DESCRIPTION: Final = ( - f"the /mcp audience only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience for LLM routes" + f"the /mcp audience or resource only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience and " + "resource for LLM routes" ) @@ -450,8 +480,12 @@ async def identity_from_id_jag( if isinstance(marker, SubjectTokenRefusal): return marker replay_key, replay_ttl_seconds = marker + mcp_resource: Final = resource_names_mcp(verified, prerequisites.base_url) + if isinstance(mcp_resource, SubjectTokenRefusal): + return mcp_resource scope: Final = granted_scope( - verified.get("scope"), mcp_resource=names_mcp_resource(verified.get("aud"), prerequisites.base_url) + verified.get("scope"), + mcp_resource=mcp_resource or names_mcp_resource(verified.get("aud"), prerequisites.base_url), ) if isinstance(scope, SubjectTokenRefusal): return scope diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b724a61cd94..96cf6ee11aa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1970,6 +1970,7 @@ class MCPServerManager: per_user_oauth_token_store: InvalidatableOAuthTokenStore | None = None, per_user_token_cache: MCPPerUserTokenCache | None = None, discovery_clock: Callable[[], float] = time.monotonic, + upstream_transport: httpx2.AsyncBaseTransport | None = None, ): self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( self.get_mcp_server_by_id @@ -1979,6 +1980,7 @@ class MCPServerManager: oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), ) + self._upstream_transport: httpx2.AsyncBaseTransport | None = upstream_transport discovery_ttl: Final = _mcp_discovery_cache_ttl() self._prompt_discovery_cache = _DiscoveryCache[ListPromptsResult]( discovery_ttl, discovery_clock, TypeAdapter(ListPromptsResult) @@ -4348,6 +4350,7 @@ class MCPServerManager: if resolved_server.allow_elicitation else None ), + upstream_transport=self._upstream_transport, ) _create_mcp_client = create_mcp_client @@ -4528,7 +4531,7 @@ class MCPServerManager: ) else: page: Final = ( - await client.list_tools_page(params) + await self._fetch_tools_page_with_timeout(client, server, params) if params is not None else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name)) ) @@ -5618,6 +5621,19 @@ class MCPServerManager: ) raise_classified_list_failure(e, server_name) + async def _fetch_tools_page_with_timeout( + self, client: MCPClient, server: MCPServer, params: PaginatedRequestParams + ) -> ListToolsResult: + deadline: Final = ( + MCP_TOOL_LISTING_TIMEOUT if server.timeout is None else max(server.timeout, MCP_TOOL_LISTING_TIMEOUT) + ) + try: + with anyio.fail_after(deadline): + return await client.list_tools_page(params) + except TimeoutError as e: + verbose_logger.warning("Timeout while listing tools from %s", server.name) + raise MCPServerListError(ServerListFault(tag="timeout"), server.name) from e + _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 def assign_unique_short_prefix( @@ -7466,7 +7482,10 @@ class MCPServerManager: return "ok" # Add timeout wrapper to prevent hanging - await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT) + await asyncio.wait_for( + client.run_with_session(_noop, quiet_on_error=True), + timeout=MCP_HEALTH_CHECK_TIMEOUT, + ) self._remember_upstream_initialize_instructions(server, client) status = "healthy" except asyncio.TimeoutError: @@ -7478,6 +7497,12 @@ class MCPServerManager: except Exception as e: health_check_error = str(e) status = "unhealthy" + verbose_logger.warning( + "MCP health check failed for %s: %s%s", + server.name, + type(e).__name__, + _upstream_failure_suffix(e), + ) return LiteLLM_MCPServerTable( server_id=server.server_id, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py index 224f956bcfd..fcd8ae22ce3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py @@ -5,7 +5,9 @@ login assertion on ID-JAG leg 1. It presents a refresh token that the IdP issued one RFC 8693 exchange (``requested_token_type`` refresh_token), which is how a SAML requesting app works at an IdP whose ID-JAG exchange does not take a ``saml2`` subject directly. That exchange has to happen at login: a SAML assertion is valid for minutes, the first MCP call that needs it can come hours later, and -the refresh token outlives the assertion by as long as the IdP decides. +the refresh token outlives the assertion by as long as the IdP decides. An OIDC login signed in through the very +IdP client the ID-JAG server uses already holds such a refresh token: it is bound to that client, so it is +presented as is and no further exchange is posted, which Okta refuses for an id_token subject anyway. The servers consulted are the two configuration authorities, the pod-local config declaration and the DB rows, like the retention gate, not the in-memory registry, which is a per-process snapshot that can lag a @@ -19,12 +21,17 @@ the ID-JAG call reports the missing refresh token as the precondition it is. from __future__ import annotations import asyncio +import os from collections.abc import Iterable, Mapping, Sequence from datetime import datetime, timezone from typing import Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + SSOClientConfig, + sso_client_config, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SSOIdentityAssertion, with_login_refresh_token, @@ -123,12 +130,31 @@ async def exchange_login_assertion_for_refresh_token( token_endpoint: TokenExchangeEndpoint, now: datetime, deadline_seconds: float = _LOGIN_EXCHANGE_DEADLINE_SECONDS, + sso_client: SSOClientConfig | None = None, ) -> SSOIdentityAssertion: """The assertion as captured when no server wants a refresh token, otherwise the assertion plus the - refresh token the first target's IdP client exchanged it for. An OIDC login's own SSO refresh token - does not count, it is bound to the SSO client and the ID-JAG client cannot redeem it.""" + refresh token one target's IdP client can redeem. An OIDC login's own refresh token is that token + when the SSO client that received it is the target's client at the same token endpoint; otherwise it + is bound to a client the ID-JAG client is not, and the assertion is exchanged for one at the first + target's client.""" if not targets or assertion.exchanged_refresh_token is not None: return assertion + reusable: Final = _target_redeeming_the_login_refresh_token(assertion, targets, sso_client) + if reusable is not None and assertion.refresh_token is not None: + verbose_proxy_logger.info( + "ID-JAG: the refresh token issued at OIDC login belongs to client_id=%s at %s, the IdP client of a server " + "with id_jag_subject_source: refresh_token, so it is kept as that subject and no further exchange is posted", + reusable.client_id, + reusable.org_token_endpoint, + ) + return with_login_refresh_token( + assertion, + assertion.refresh_token.get_secret_value(), + None, + now, + endpoint=reusable.org_token_endpoint, + client_id=reusable.client_id, + ) config: Final = targets[0] if len(targets) > 1: verbose_proxy_logger.warning( @@ -185,6 +211,24 @@ async def exchange_login_assertion_for_refresh_token( ) +def _target_redeeming_the_login_refresh_token( + assertion: SSOIdentityAssertion, targets: tuple[IdJagConfig, ...], sso_client: SSOClientConfig | None +) -> IdJagConfig | None: + """The target whose IdP client is the SSO client the login's refresh token was issued to. A refresh + token is redeemable only by the client that received it at the endpoint that issued it, so both have + to match exactly before the login token is presented anywhere.""" + if assertion.refresh_token is None or sso_client is None: + return None + return next( + ( + config + for config in targets + if config.org_token_endpoint == sso_client.token_endpoint and config.client_id == sso_client.client_id + ), + None, + ) + + async def exchange_login_assertion_if_configured( assertion: SSOIdentityAssertion | None, ) -> SSOIdentityAssertion | None: @@ -209,7 +253,11 @@ async def exchange_login_assertion_if_configured( servers: Final = await login_exchange_servers(global_mcp_server_manager, db_rows) targets: Final = refresh_token_exchange_targets(to_server_spec(server) for server in servers) return await exchange_login_assertion_for_refresh_token( - assertion, targets, TokenEndpointClient(), datetime.now(timezone.utc) + assertion, + targets, + TokenEndpointClient(), + datetime.now(timezone.utc), + sso_client=sso_client_config(os.environ), ) except Exception as exc: # noqa: BLE001 # the login itself must not fail on an egress-side exchange verbose_proxy_logger.warning("ID-JAG: the login-time refresh token exchange failed unexpectedly: %s", exc) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index bd0813dd966..ea0cdee5c1c 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -429,10 +429,11 @@ def _login_refresh_token_subject( assertion: SSOIdentityAssertion, config: IdJagConfig, now: datetime ) -> Result[_IdJagSubjectToken, CredError]: """The subject for a server configured with ``id_jag_subject_source: refresh_token``: the refresh - token the ID-JAG client exchanged the login assertion for when the user signed in. Never the SSO - refresh token an OIDC login holds, that one is bound to the SSO client. The assertion's own expiry - does not apply here, a SAML assertion is spent minutes after login while the refresh token it - bought lives as long as the IdP said.""" + token the ID-JAG client holds from the user's sign-in, exchanged from the login assertion or, when + the SSO client is this server's IdP client, the one the OIDC login itself returned. The login-time + hook records which client and endpoint it is bound to, and a token bound elsewhere is never sent + here. The assertion's own expiry does not apply, a SAML assertion is spent minutes after login + while the refresh token it bought lives as long as the IdP said.""" if assertion.exchanged_refresh_token is None: return Error( CredError.of_precondition_required( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 58494654082..19b3bc294d5 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -229,10 +229,19 @@ def carry_exchanged_refresh_token( ) -> SSOIdentityAssertion: """The renewed OIDC assertion still carrying the refresh token the ID-JAG client exchanged at login. Renewal only replaces the id_token and the SSO refresh token, the exchanged one is unrelated to - them and stays valid for as long as the IdP said.""" + them and stays valid for as long as the IdP said. When the login's own refresh token was kept as + that subject (the SSO client is the ID-JAG client), a rotation the renewal just received replaces + it too, since the IdP has retired the token the row held.""" + login_token_reused: Final = ( + previous.refresh_token is not None + and previous.exchanged_refresh_token is not None + and previous.exchanged_refresh_token.get_secret_value() == previous.refresh_token.get_secret_value() + ) return renewed.model_copy( update={ - "exchanged_refresh_token": previous.exchanged_refresh_token, + "exchanged_refresh_token": renewed.refresh_token + if login_token_reused + else previous.exchanged_refresh_token, "exchanged_refresh_token_expires_at": previous.exchanged_refresh_token_expires_at, "exchanged_refresh_token_endpoint": previous.exchanged_refresh_token_endpoint, "exchanged_refresh_token_client_id": previous.exchanged_refresh_token_client_id, diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index b19f904d81d..3b7cf904b6a 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -207,6 +207,7 @@ async def prepare_upstream_client( protocol_version: MCPUpstreamProtocol, sampling_callback: SamplingCallback | None = None, elicitation_callback: ElicitationCallback | None = None, + upstream_transport: httpx2.AsyncBaseTransport | None = None, ) -> MCPClient: transport: Final[MCPTransportType] = server.transport or MCPTransport.sse server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server) @@ -263,6 +264,7 @@ async def prepare_upstream_client( server_url=server.url or "", transport_type=transport, protocol_version=protocol_version, + upstream_transport=upstream_transport, auth_type=server.auth_type, timeout=timeout, extra_headers=resolved_headers, @@ -304,6 +306,7 @@ async def prepare_upstream_client( server_url=server.url or "", transport_type=transport, protocol_version=protocol_version, + upstream_transport=upstream_transport, auth_type=server.auth_type, auth_value=auth_value, auth_header_name=auth_header_name, diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 8dab693cd58..f5e8304c83c 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -72,14 +72,22 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spe _LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp")) _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS)) -# Method-scoped because the generic credential CRUD handlers share these paths: -# the route check cannot admit PATCH/DELETE on a name like "user_connections" -# or "*/user_connection*" without also opening the admin CRUD handlers. -_INTERNAL_USER_CREDENTIAL_CONNECTION_ROUTES: Final = ( +# Writes that only touch the caller's own credential row (BYOK connections and +# per-user MCP credentials). Method-scoped because the generic credential CRUD +# handlers share these paths: the route check cannot admit PATCH/DELETE on a +# name like "user_connections" without also opening the admin CRUD handlers. +_SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( ("GET", "/credentials/user_connections"), ("POST", "/credentials/{credential_name:path}/user_connection/start"), ("POST", "/credentials/{credential_name:path}/user_connection/poll"), ("DELETE", "/credentials/{credential_name:path}/user_connection"), + ("POST", "/v1/mcp/server/oauth/{server_id}/register"), + ("POST", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-env-vars"), + ("DELETE", "/v1/mcp/server/{server_id}/user-env-vars"), ) @@ -349,7 +357,7 @@ class RouteChecks: _user_role == LitellmUserRoles.INTERNAL_USER.value and ( RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value) - or RouteChecks.is_internal_user_credential_connection_route(route=route, request=request) + or RouteChecks.is_self_scoped_credential_route(route=route, request=request) ) or user_is_org_admin(request_data=request_data, user_object=user_obj) and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value) @@ -702,13 +710,13 @@ class RouteChecks: _get_request_method = get_request_method @staticmethod - def is_internal_user_credential_connection_route(route: str, request: Request | None) -> bool: + def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: method: Final = RouteChecks.get_request_method(request) if method is None: return False return any( method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,)) - for expected_method, pattern in _INTERNAL_USER_CREDENTIAL_CONNECTION_ROUTES + for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES ) @staticmethod @@ -1033,6 +1041,9 @@ class RouteChecks: if route == "/session/logout": return + if RouteChecks.is_self_scoped_credential_route(route=route, request=request): + return + # Hard-block known write routes regardless of HTTP method (defensive # — these are POSTs in practice, but pinning them here protects # against future GET-shaped writes). diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 5fa52e5e4f4..03540a7395e 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -22,7 +22,7 @@ from litellm.proxy.common_utils.codex_model_catalog import ( CodexStockModel, codex_models_response_json, ) -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel from .auth import cli_context_obj, context_secret_vault, get_stored_api_key, login from .claude_settings import ClaudeSettingsError, install_statusline_script diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py index 28a57f71d08..eaee886f148 100644 --- a/litellm/proxy/client/cli/commands/autoroute/config.py +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -3,7 +3,7 @@ from typing import Final, Literal from pydantic import ConfigDict, Field, JsonValue, TypeAdapter -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel TIER_NAMES: Final[tuple[str, ...]] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") AUTOROUTER_MODEL_NAME: Final = "autorouter" diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index cdceb91712b..60492a81cab 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -35,7 +35,7 @@ from litellm.litellm_core_utils.private_json import ( stage_private_json, write_private_bytes, ) -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel from . import statusline_script from .cmd_quoting import quote_for_cmd diff --git a/litellm/proxy/client/cli/commands/codex_settings.py b/litellm/proxy/client/cli/commands/codex_settings.py index ca133d3f1e5..557f4379773 100644 --- a/litellm/proxy/client/cli/commands/codex_settings.py +++ b/litellm/proxy/client/cli/commands/codex_settings.py @@ -23,7 +23,7 @@ from litellm.litellm_core_utils.private_json import ( stage_private_bytes, stage_private_json, ) -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel from .agents import CODEX_PROXY_PROVIDER, codex_proxy_provider diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py index 45c79ccf014..7903a978a4a 100644 --- a/litellm/proxy/client/cli/commands/configure.py +++ b/litellm/proxy/client/cli/commands/configure.py @@ -7,7 +7,7 @@ from typing import Final import click from InquirerPy import inquirer -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel from .auth import CliContextObj from .claude_settings import ( diff --git a/litellm/proxy/client/cli/commands/configure_profiles.py b/litellm/proxy/client/cli/commands/configure_profiles.py index 3ae651ff6db..dfd855dcb9b 100644 --- a/litellm/proxy/client/cli/commands/configure_profiles.py +++ b/litellm/proxy/client/cli/commands/configure_profiles.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.private_json import ( ensure_private_dir, stage_private_json, ) -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel from .agents import codex_config_path from .claude_settings import claude_settings_path, configure_state_path diff --git a/litellm/proxy/client/cli/commands/configure_setup.py b/litellm/proxy/client/cli/commands/configure_setup.py index 7a262650cc2..6004de72e18 100644 --- a/litellm/proxy/client/cli/commands/configure_setup.py +++ b/litellm/proxy/client/cli/commands/configure_setup.py @@ -14,12 +14,8 @@ from InquirerPy import inquirer from InquirerPy.base.control import Choice from pydantic import TypeAdapter, ValidationError -from litellm.proxy.common_utils.model_listing_utils import ( - CLAUDE_CODE_CLIENT, - CLAUDE_CODE_PICKER_PATTERN, - GATEWAY_CLIENT_HEADER, -) -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.constants import CLAUDE_CODE_CLIENT, CLAUDE_CODE_PICKER_PATTERN, GATEWAY_CLIENT_HEADER +from litellm.types.llms.base_model import LiteLLMBaseModel from .agents import codex_config_path from .claude_settings import ( diff --git a/litellm/proxy/client/cli/commands/debug.py b/litellm/proxy/client/cli/commands/debug.py index 3db3662b021..8728a9323b7 100644 --- a/litellm/proxy/client/cli/commands/debug.py +++ b/litellm/proxy/client/cli/commands/debug.py @@ -20,7 +20,7 @@ import click import requests from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel from ...http_client import HTTPClient from ._cli_context import cli_context_values diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 4cb2372e07e..b228617699c 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -19,7 +19,7 @@ import requests from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator from pydantic.types import StringConstraints -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" PI_PROVIDER_NAME: Final = "litellm" diff --git a/litellm/proxy/client/cli/commands/pkce_login.py b/litellm/proxy/client/cli/commands/pkce_login.py index 7cf846bdf9a..c9dd542eb36 100644 --- a/litellm/proxy/client/cli/commands/pkce_login.py +++ b/litellm/proxy/client/cli/commands/pkce_login.py @@ -24,7 +24,7 @@ from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.cli_token_utils import CLI_TOKEN_FRESHNESS_BUFFER_SECONDS -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel if TYPE_CHECKING: from .auth import CliTokenData diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index f6ecaf1a9be..f7c545ef48e 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -8,6 +8,7 @@ from typing import Final, TypeAlias import click from litellm._logging import verbose_logger +from litellm.proxy.common_utils.banner import show_banner from .commands.auth import cli_context_obj, context_secret_vault @@ -115,8 +116,6 @@ def show_commands(): def setup_shell(ctx: click.Context): """Set up the interactive shell with banner and initial info.""" - from litellm.proxy.common_utils.banner import show_banner - show_banner() # Show server connection info @@ -137,8 +136,6 @@ def handle_special_commands(user_input: str) -> bool: return True elif user_input.lower() == "clear": click.clear() - from litellm.proxy.common_utils.banner import show_banner - show_banner() show_commands() return True diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 0f4bbc911c6..3f1e7f4259c 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -1,42 +1,20 @@ -# stdlib imports - -# third party imports +import functools +import importlib +from collections.abc import Callable, Iterator, Mapping, MutableMapping from typing import Final, Literal, Protocol import click from litellm._version import version as litellm_version from litellm.litellm_core_utils.cli_token_utils import load_cli_token -from litellm.proxy.client.health import HealthManagementClient -from .commands.agents import agent_commands from .commands.auth import ( - auth_group, cli_context_obj, context_secret_vault, get_stored_api_key, keychain_unreadable_notice, - login, - logout, - whoami, ) -from .commands.autoroute.commands import autoroute_group -from .commands.chat import chat -from .commands.config import config_commands, get_config_value, hidden_command_names -from .commands.configure import configure_group, reconfigure_group, unconfigure_group -from .commands.credentials import credentials -from .commands.debug import debug -from .commands.encryption import encryption -from .commands.http import http -from .commands.keys import keys -from .commands.model_groups import model_groups - -# local imports -from .commands.models import models -from .commands.teams import teams -from .commands.up import down, up -from .commands.users import users -from .interface import interactive_shell +from .commands.config import get_config_value, hidden_command_names class _CliContextObjSlots(Protocol): @@ -62,6 +40,8 @@ def print_version(base_url: str, api_key: str | None): click.echo(f"LiteLLM Proxy CLI Version: {litellm_version}") if base_url: click.echo(f"LiteLLM Proxy Server URL: {base_url}") + from litellm.proxy.client.health import HealthManagementClient + try: health_client: Final = HealthManagementClient(base_url=base_url, api_key=api_key) server_version: Final = health_client.get_server_version() @@ -87,7 +67,81 @@ class HideConfiguredCommandsGroup(click.Group): return [name for name in super().list_commands(ctx) if name not in hidden] -@click.group(cls=HideConfiguredCommandsGroup, invoke_without_command=True) +class LazyCommands(MutableMapping[str, click.Command]): + """Subcommands by name, each imported only when Click looks it up, so `lite ` loads just that command""" + + def __init__(self, loaders: Mapping[str, Callable[[], click.Command]]) -> None: + self._loaders = dict(loaders) + + def __getitem__(self, name: str) -> click.Command: + return self._loaders[name]() + + def __setitem__(self, name: str, command: click.Command) -> None: + self._loaders[name] = lambda: command + + def __delitem__(self, name: str) -> None: + del self._loaders[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._loaders) + + def __len__(self) -> int: + return len(self._loaders) + + +def _command_in(module: str, attribute: str) -> Callable[[], click.Command]: + @functools.cache + def load() -> click.Command: + match getattr(importlib.import_module(module, __package__), attribute): + case click.Command() as command: + return command + case _: + raise TypeError(f"{module}.{attribute} is not a click command") + + return load + + +def _agent_command(name: str) -> Callable[[], click.Command]: + @functools.cache + def load() -> click.Command: + from .commands.agents import agent_commands + + return next(command for command in agent_commands() if command.name == name) + + return load + + +AGENT_COMMAND_NAMES: Final = ("claude", "codex", "opencode", "pi") + +LITE_COMMANDS: Final = LazyCommands( + { + "login": _command_in(".commands.auth", "login"), + "logout": _command_in(".commands.auth", "logout"), + "whoami": _command_in(".commands.auth", "whoami"), + "auth": _command_in(".commands.auth", "auth_group"), + "models": _command_in(".commands.models", "models"), + "credentials": _command_in(".commands.credentials", "credentials"), + "encryption": _command_in(".commands.encryption", "encryption"), + "chat": _command_in(".commands.chat", "chat"), + "http": _command_in(".commands.http", "http"), + "debug": _command_in(".commands.debug", "debug"), + "keys": _command_in(".commands.keys", "keys"), + "teams": _command_in(".commands.teams", "teams"), + "users": _command_in(".commands.users", "users"), + **{name: _agent_command(name) for name in AGENT_COMMAND_NAMES}, + "up": _command_in(".commands.up", "up"), + "down": _command_in(".commands.up", "down"), + "model-groups": _command_in(".commands.model_groups", "model_groups"), + "autoroute": _command_in(".commands.autoroute.commands", "autoroute_group"), + "config": _command_in(".commands.config", "config_commands"), + "configure": _command_in(".commands.configure", "configure_group"), + "reconfigure": _command_in(".commands.configure", "reconfigure_group"), + "unconfigure": _command_in(".commands.configure", "unconfigure_group"), + } +) + + +@click.group(cls=HideConfiguredCommandsGroup, invoke_without_command=True, commands=LITE_COMMANDS) @click.option( "--version", "-v", @@ -166,6 +220,8 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s # If no subcommand was invoked, start interactive mode if ctx.invoked_subcommand is None: + from .interface import interactive_shell + interactive_shell(ctx, cli) @@ -177,46 +233,6 @@ def version(ctx: click.Context): print_version(ctx_obj.get("base_url"), ctx_obj.get("api_key")) -# Add authentication commands as top-level commands -cli.add_command(login) -cli.add_command(logout) -cli.add_command(whoami) -# Add the auth command group (e.g. `lite auth print-token`, used as Claude Code's apiKeyHelper) -cli.add_command(auth_group, name="auth") -# Add the models command group -cli.add_command(models) -# Add the credentials command group -cli.add_command(credentials) -# Add the encryption migration command group -cli.add_command(encryption) -# Add the chat command group -cli.add_command(chat) -# Add the http command group -cli.add_command(http) -cli.add_command(debug) -# Add the keys command group -cli.add_command(keys) -# Add the teams command group -cli.add_command(teams) -# Add the users command group -cli.add_command(users) -# Add a top-level command per coding agent (claude, codex, opencode, ...) -for agent_command in agent_commands(): - cli.add_command(agent_command) -# Add the up/down commands (route Claude Code through the local LiteLLM proxy) -cli.add_command(up) -cli.add_command(down) -# Add the model-groups command group (discover models your key can access) -cli.add_command(model_groups) -# Add the autoroute command group (QA auto-routing against your real proxy) -cli.add_command(autoroute_group, name="autoroute") -cli.add_command(config_commands) -# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key) -cli.add_command(configure_group) -cli.add_command(reconfigure_group) -cli.add_command(unconfigure_group) - - LITELLM_PROXY_DEPRECATION_NOTICE: Final = ( "The `litellm-proxy` command is deprecated and will be removed in a future release; " "run `lite` instead, it takes the same commands and options." diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2d387b43aa8..e52f11641b0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,7 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -138,6 +138,7 @@ from litellm.types.router_weights import validate_router_weights _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") +_ChunkT: Final = TypeVar("_ChunkT") KNOWN_PROXY_ROUTES: Final = frozenset( route for member in LiteLLMRoutes for route in member.value if route.startswith("/") @@ -506,6 +507,10 @@ async def _cancel_pending_gather_tasks(tasks: Sequence["asyncio.Task[object]"]) pass +def _encodes_as_latin1(value: str) -> bool: + return all(ord(character) < 256 for character in value) + + @lru_cache(maxsize=512) def _litellm_model_supports_stream_options(litellm_model: str) -> bool: try: @@ -923,6 +928,21 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict: return default_error +async def _close_upstream(upstream: object) -> None: + aclose: Final[Callable[[], Awaitable[object]] | None] = getattr(upstream, "aclose", None) + if aclose is not None: + await aclose() + + +async def _close_upstream_and_log_its_spend(upstream: object) -> bool: + """Close the upstream and wait until the spend it collected has been logged. False when it had none to log.""" + await _close_upstream(upstream) + wait_for_spend_flush: Final[Callable[[], Awaitable[object]] | None] = getattr( + upstream, "wait_for_spend_flush", None + ) + return wait_for_spend_flush is not None and await wait_for_spend_flush() is True + + class _UpstreamClosingStreamingResponse(StreamingResponse): """StreamingResponse that always closes its body iterator and the wrapped upstream generator. @@ -969,7 +989,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse): class _ClientDisconnectedBeforeFirstChunk(Exception): - """Client went away during create_response's first-chunk buffering window. + """Client went away while the first chunk of a streamed response was being buffered. The upstream LLM stream has already been closed by the time this is raised. """ @@ -993,24 +1013,24 @@ async def _wait_for_http_disconnect(request: Request) -> None: async def _buffer_first_chunk_honoring_disconnect( - generator: AsyncGenerator[str, None], + generator: AsyncGenerator[_ChunkT, None], request: Request | None, -) -> str: +) -> _ChunkT: """Fetch the first streamed chunk, cancelling the upstream LLM call if the client disconnects before it arrives. - create_response buffers the first chunk to detect error-only streams before - handing the StreamingResponse to Starlette, which only begins listening for - client disconnects once it is serving that response. A disconnect during a - long time-to-first-token would otherwise leave the upstream call running - until the request timeout (LIT-3568). Cancelling the fetch propagates into - async_streaming_data_generator, whose finally block records the 499 and - closes the upstream stream. + create_response and the guarded passthrough relay buffer the first chunk + before handing the StreamingResponse to Starlette, which only begins + listening for client disconnects once it is serving that response. A + disconnect during a long time-to-first-token would otherwise leave the + upstream call running until the request timeout (LIT-3568). Cancelling the + fetch propagates into the generator, whose cleanup closes the upstream + stream. A cancellation of the caller stops the fetch the same way. """ if request is None: return await generator.__anext__() - chunk_task: Final[asyncio.Task[str]] = asyncio.ensure_future(generator.__anext__()) + chunk_task: Final[asyncio.Task[_ChunkT]] = asyncio.ensure_future(generator.__anext__()) disconnect_task: Final[asyncio.Task[None]] = asyncio.ensure_future(_wait_for_http_disconnect(request)) try: await asyncio.wait({chunk_task, disconnect_task}, return_when=asyncio.FIRST_COMPLETED) @@ -1019,6 +1039,9 @@ async def _buffer_first_chunk_honoring_disconnect( # Take the cancellation path whenever a disconnect was observed, even if # the first chunk landed in the same scheduler turn. disconnect_observed: Final = disconnect_task.done() + except BaseException: + await _cancel_first_chunk_fetch(chunk_task, generator) + raise finally: disconnect_task.cancel() try: @@ -1029,6 +1052,14 @@ async def _buffer_first_chunk_honoring_disconnect( if not disconnect_observed and chunk_task.done() and not chunk_task.cancelled(): return chunk_task.result() + await _cancel_first_chunk_fetch(chunk_task, generator) + verbose_proxy_logger.info("client disconnected before the first streamed chunk, upstream LLM request cancelled") + raise _ClientDisconnectedBeforeFirstChunk() + + +async def _cancel_first_chunk_fetch( + chunk_task: asyncio.Task[_ChunkT], generator: AsyncGenerator[_ChunkT, None] +) -> None: chunk_task.cancel() with anyio.CancelScope(shield=True): try: @@ -1038,9 +1069,25 @@ async def _buffer_first_chunk_honoring_disconnect( try: await generator.aclose() except BaseException as exc: # noqa: BLE001 - verbose_proxy_logger.debug("create_response: error closing generator on disconnect: %s", exc) - verbose_proxy_logger.info("create_response: client disconnected before first chunk, upstream LLM request cancelled") - raise _ClientDisconnectedBeforeFirstChunk() + verbose_proxy_logger.debug( + "error closing the streamed response's generator before its first chunk: %s", exc + ) + + +def _client_disconnected_response(headers: Mapping[str, str] | None) -> JSONResponse: + """The 499 returned once the client left before the first streamed chunk, which the client never reads.""" + return JSONResponse( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, + content={ + "error": { + "message": _CLIENT_DISCONNECT_DETAIL, + "type": "client_disconnect", + "param": None, + "code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED), + } + }, + headers=headers, + ) def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: @@ -1170,18 +1217,7 @@ async def create_response( except _ClientDisconnectedBeforeFirstChunk: # Client vanished during the time-to-first-token wait; the upstream # stream is already closed. Return a 499 the (now-gone) client never reads. - return JSONResponse( - status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, - content={ - "error": { - "message": _CLIENT_DISCONNECT_DETAIL, - "type": "client_disconnect", - "param": None, - "code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED), - } - }, - headers=headers, - ) + return _client_disconnected_response(headers) except StopAsyncIteration: # Generator was empty. Default status async def empty_gen() -> AsyncGenerator[str, None]: @@ -1627,19 +1663,26 @@ def _totals_to_zero(response_cost: float | str | None) -> bool: return False +def _is_cache_hit(litellm_logging_obj: LiteLLMLoggingObj) -> bool: + caching_details: Final = litellm_logging_obj.caching_details + return caching_details is not None and caching_details.get("cache_hit") is True + + def _get_cost_breakdown_from_logging_obj( litellm_logging_obj: LiteLLMLoggingObj | None, response_cost: float | str | None = None, ) -> CostBreakdownHeaderValues: """Extract discount, margin, and per-component cost information from logging object's cost breakdown. - A non-inference call that priced at zero never records a breakdown, so its components are - reported as zero here. Any such call that did price normally (retrieving a background response, - and the cost poller's read of one) reports the breakdown it stored, or nothing at all when the - breakdown has not landed yet. + A cache hit reports every component as zero. A non-inference call that priced at zero never + records a breakdown, so its components are reported as zero too. Any such call that did price + normally (retrieving a background response, and the cost poller's read of one) reports the + breakdown it stored, or nothing at all when the breakdown has not landed yet. """ if not litellm_logging_obj or not hasattr(litellm_logging_obj, "cost_breakdown"): return CostBreakdownHeaderValues() + if _is_cache_hit(litellm_logging_obj): + return _ZERO_COST_BREAKDOWN cost_breakdown: Final = litellm_logging_obj.cost_breakdown if not cost_breakdown: @@ -1669,6 +1712,13 @@ def _header_value(value: float | None) -> str | None: return str(value) if value is not None else None +def _response_cost_header_value(hidden_params: Mapping[str, object]) -> float | str: + if hidden_params.get("cache_hit") is True: + return 0.0 + response_cost: Final = hidden_params.get("response_cost") + return response_cost if isinstance(response_cost, (int, float, str)) and response_cost else "" + + def cost_headers( response_cost: float | str | None, litellm_logging_obj: LiteLLMLoggingObj | None ) -> Mapping[str, str | None]: @@ -2027,10 +2077,11 @@ class ProxyBaseLLMRequestProcessing: headers.update(logging_caching_headers) try: - return {key: str(value) for key, value in headers.items() if value not in exclude_values} + present: Final = {key: str(value) for key, value in headers.items() if value not in exclude_values} except Exception as e: verbose_proxy_logger.error("Error setting custom headers: %s", e) return {} + return {key: value for key, value in present.items() if _encodes_as_latin1(value)} @staticmethod async def build_litellm_proxy_success_headers_from_llm_response( @@ -2047,18 +2098,12 @@ class ProxyBaseLLMRequestProcessing: Build LiteLLM proxy response headers for routes that call the LLM directly (e.g. Google native :generateContent) instead of base_process_llm_request. """ - if isinstance(response, dict): - hidden_params = get_hidden_params_dict(response) - else: - hidden_params = getattr(response, "_hidden_params", None) or {} - if not isinstance(hidden_params, dict): - hidden_params = {} - + hidden_params: Final = get_hidden_params_dict(response) model_id: Final = ProxyBaseLLMRequestProcessing._get_model_id_from_response(hidden_params, request_data) cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" - response_cost: Final = hidden_params.get("response_cost", None) or "" + response_cost: Final = _response_cost_header_value(hidden_params) fastest_response_batch_completion: Final = hidden_params.get("fastest_response_batch_completion", None) additional_headers: Final = hidden_params.get("additional_headers", {}) or {} @@ -2603,7 +2648,7 @@ class ProxyBaseLLMRequestProcessing: cache_key=hidden_params.get("cache_key") or "", api_base=hidden_params.get("api_base") or "", version=version, - response_cost=hidden_params.get("response_cost") or "", + response_cost=_response_cost_header_value(hidden_params), model_region=getattr(user_api_key_dict, "allowed_model_region", ""), fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"), request_data=self.data, @@ -2648,13 +2693,17 @@ class ProxyBaseLLMRequestProcessing: hold the attribute at all, and Google :generateContent carries ``_hidden_params`` but no synchronously-populated ``response_cost``. In both cases the cost is read back from the logging object instead, recomputing from - the same calculator only when it has not been stored yet. + the same calculator only when it has not been stored yet. A response served + from the cache costs 0, which the calculator cannot tell from a fresh call, and any + other zero or unknown cost comes back empty so no cost header is sent. """ + if _is_cache_hit(logging_obj): + return 0.0 stored_cost: Final = logging_obj.model_call_details.get("response_cost") if isinstance(stored_cost, (int, float)): - return float(stored_cost) + return float(stored_cost) or "" recomputed_cost: Final = logging_obj.response_cost_calculator(result=response) - return recomputed_cost if isinstance(recomputed_cost, (int, float)) else "" + return recomputed_cost or "" if isinstance(recomputed_cost, (int, float)) else "" def _debug_log_request_payload(self) -> None: """Log request payload at DEBUG level, truncating if too large.""" @@ -2863,7 +2912,7 @@ class ProxyBaseLLMRequestProcessing: cache_key, api_base, response_cost = ( hidden_params.get("cache_key", None) or "", hidden_params.get("api_base", None) or "", - hidden_params.get("response_cost", None) or "", + _response_cost_header_value(hidden_params), ) fastest_response_batch_completion, additional_headers = ( hidden_params.get("fastest_response_batch_completion", None), @@ -2945,10 +2994,25 @@ class ProxyBaseLLMRequestProcessing: else: generator = response - if ( - self._has_post_call_guardrails_for_passthrough() - and self._passthrough_endpoint_has_stream_guardrail_handler() - ): + post_call_guarded: Final = self._has_post_call_guardrails_for_passthrough() + relayed_guardrails: Final = ( + self._passthrough_stream_guardrails( + self.data # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + ) + if post_call_guarded and self._passthrough_stream_is_relayable() + else None + ) + if relayed_guardrails is not None: + return await self._relay_guarded_passthrough_event_stream( + request=request, + request_data=self.data, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + generator=generator, # pyright: ignore[reportUnknownArgumentType] # raw upstream byte stream + guardrails=relayed_guardrails, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, + headers=streaming_headers, # pyright: ignore[reportUnknownArgumentType] # merged upstream headers + ) + if post_call_guarded and self._passthrough_endpoint_has_stream_guardrail_handler(): body_bytes: Final = b"".join([chunk async for chunk in generator]) modified_bytes: Final = await self._handle_event_stream_allm_passthrough_route( body_bytes=body_bytes, @@ -3118,7 +3182,7 @@ class ProxyBaseLLMRequestProcessing: recover_response_cost: Final = not response_cost and hidden_params.get("response_cost") is None computed_cost_for_headers: Final = ( - self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) or "" + self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) if recover_response_cost else response_cost ) @@ -3456,16 +3520,199 @@ class ProxyBaseLLMRequestProcessing: endpoints may have their stream buffered for post-call guardrails; every other endpoint must keep streaming so the response is not silently turned into a non-streaming body when no content modification would occur (e.g. - Bedrock invoke-with-response-stream, whose frames the Converse handler - leaves untouched). + Bedrock invoke-with-response-stream of a model other than Claude). """ from litellm.llms.pass_through.guardrail_translation.handler import ( LlmPassthroughRouteHandler, ) return LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - self.data.get("custom_llm_provider"), - self.data.get("endpoint"), + self.data.get("custom_llm_provider"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("endpoint"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + getattr(self.data.get("litellm_logging_obj"), "model", None), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("data"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + ) + + def _passthrough_stream_is_relayable(self) -> bool: + """True when the passthrough provider's guardrail translation scans this endpoint's stream as it relays.""" + from litellm.llms.pass_through.guardrail_translation.handler import ( + LlmPassthroughRouteHandler, + ) + + return LlmPassthroughRouteHandler.supports_event_stream_relay( + self.data.get("custom_llm_provider"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("endpoint"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + getattr(self.data.get("litellm_logging_obj"), "model", None), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("data"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + ) + + @staticmethod + def _passthrough_stream_guardrails( + request_data: dict[str, object], # mutable-ok: model-level guardrails merge into a copy of it + ) -> tuple[tuple[CustomGuardrail, bool], ...] | None: + """ + The post_call guardrails to run on a passthrough event stream relayed frame by frame, each with whether it + holds the stream until moderated by default, or None when the stream must be buffered whole instead. + + A guardrail relays when ``post_call_success_hook`` would scan it through ``apply_guardrail``. One that runs + its own ``async_post_call_success_hook``, masks responses with ``mask_response_content``, releases a held + stream as each scan passes, runs with ``run_in_parallel`` after the sequential guardrails, or is a step of a + post_call policy pipeline still needs the whole response, so its stream stays buffered. + """ + from litellm.proxy.proxy_server import llm_router + from litellm.proxy.utils import pipeline_managed_guardrail_names + + if pipeline_managed_guardrail_names(request_data, "post_call"): + return None + guardrail_data: Final = check_and_merge_model_level_guardrails( # pyright: ignore[reportUnknownVariableType] # returns the untyped request dict + data=request_data, llm_router=llm_router + ) + relayed: Final[list[tuple[CustomGuardrail, bool]]] = [] # mutable-ok: built once per request + for cb in litellm.callbacks: + if not isinstance(cb, CustomGuardrail): + continue + if ( + cb.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + data=guardrail_data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + continue + leaf_attrs = type(cb).__dict__ + holds_until_moderated = "async_post_call_streaming_iterator_hook" in leaf_attrs + if ( + "apply_guardrail" not in leaf_attrs + or cb.use_native_lifecycle_hooks + or cb.mask_response_content + or cb.run_in_parallel + or ProxyBaseLLMRequestProcessing._releases_held_stream_on_scan(cb, holds_until_moderated) + ): + return None + relayed.append((cb, holds_until_moderated)) + return tuple(relayed) + + @staticmethod + def _releases_held_stream_on_scan(guardrail: CustomGuardrail, holds_until_moderated: bool) -> bool: + """ + Whether the unified guardrail holds this guardrail's stream and releases each window once its scan passes + (``streaming_buffer_release_on_scan``). Those windows go out as received, without the scan's masks. + """ + from litellm.proxy.utils import unified_guardrail + + def flag(name: str, default: bool) -> bool: + return bool(unified_guardrail.resolve_streaming_flag(guardrail, name, default)) + + return ( + flag("streaming_buffer_until_moderated", holds_until_moderated) + and flag("streaming_buffer_release_on_scan", False) + and not flag("streaming_end_of_stream_only", False) + ) + + async def _relay_guarded_passthrough_event_stream( + self, + request: Request | None, + request_data: dict[str, object], # mutable-ok: the request dict the hooks mutate + generator: AsyncIterator[bytes], + guardrails: Sequence[tuple[CustomGuardrail, bool]], + proxy_logging_obj: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, + headers: Mapping[str, str], + ) -> StreamingResponse | JSONResponse: + """ + Relay a passthrough event stream as it arrives while post_call guardrails scan it, the way /v1/messages + streams are guarded. A failure before the first frame raises, so the caller gets an HTTP error as on the + buffered path; after it, the failure ends the stream with the provider's in-stream exception frame. A + client that leaves before the first frame stops the upstream read and gets the 499 chat completions returns. + """ + from litellm.llms.pass_through.guardrail_translation.handler import ( + LlmPassthroughRouteHandler, + ) + + provider: Final = request_data.get("custom_llm_provider") + items: Final = LlmPassthroughRouteHandler.event_stream_items(provider, request_data.get("endpoint"), generator) + if items is None: + raise TypeError(f"provider {provider!r} has no event-stream framer") + guarded: Final = proxy_logging_obj.async_post_call_event_stream_iterator_hook( + response=items, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + guardrails=guardrails, + ) + + async def relay() -> AsyncGenerator[bytes, None]: + delivered = False # rebind-ok: set once the first frame reached the caller + try: + async for item in guarded: + frame = LlmPassthroughRouteHandler.encode_event_stream_item(provider, item) + if frame: + yield frame + delivered = True + except Exception as e: + if not delivered: + raise + verbose_proxy_logger.exception("allm_passthrough_route: guarded event stream failed mid-stream: %s", e) + # The failure hook rewrites the request's logging object, so the stream's spend is logged first. + await _close_upstream_and_log_its_spend(generator) + transformed: Final = await proxy_logging_obj.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_data, + ) + failure: Final = transformed if transformed is not None else e + message: Final = ( + serialize_http_exception_detail(failure.detail)[0] + if isinstance(failure, HTTPException) + else redact_internal_details_from_client_message(getattr(failure, "message", str(failure))) + ) + error_item: Final = LlmPassthroughRouteHandler.event_stream_error_item( + provider, message, error_status_code(failure, status.HTTP_500_INTERNAL_SERVER_ERROR) + ) + if error_item is None: + raise + yield error_item + finally: + try: + await guarded.aclose() + finally: + await _close_upstream(generator) + + relayed: Final = relay() + read_ahead: Final[list[bytes]] = [] # mutable-ok: the first frame, read before the response starts + try: + read_ahead.append(await _buffer_first_chunk_honoring_disconnect(relayed, request)) + except _ClientDisconnectedBeforeFirstChunk: + with anyio.CancelScope(shield=True): + if not await _close_upstream_and_log_its_spend(generator): + # No spend row settles a stream that returned nothing it could log, so release what the request + # holds, as chat completions does. + from litellm.proxy.spend_tracking.budget_reservation import ( + release_budget_reservation_on_cancel, # pyright: ignore[reportUnknownVariableType] # takes the untyped reservation dict + ) + + await release_budget_reservation_on_cancel(getattr(user_api_key_dict, "budget_reservation", None)) + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + return _client_disconnected_response( + {k: v for k, v in headers.items() if k.lower() not in ("content-length", "content-type")} + ) + except StopAsyncIteration: + pass + except BaseException: + with anyio.CancelScope(shield=True): + await _close_upstream_and_log_its_spend(generator) + raise + + async def body() -> AsyncGenerator[bytes, None]: + async with contextlib.aclosing(relayed) as rest: + for frame in read_ahead: + yield frame + async for frame in rest: + yield frame + + return _UpstreamClosingStreamingResponse( + content=body(), # pyright: ignore[reportArgumentType] # bytes body, as on the unguarded relay + media_type=self._passthrough_event_stream_media_type(), + headers={k: v for k, v in headers.items() if k.lower() != "content-length"}, + upstream_generator=relayed, # pyright: ignore[reportArgumentType] # closing it closes the guardrails and the upstream ) def _passthrough_event_stream_media_type(self) -> str | None: diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 825b07b3c3c..9b42f99305e 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -29,7 +29,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal from pydantic import ConfigDict, StringConstraints, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.base_model import LiteLLMBaseModel if TYPE_CHECKING: from litellm.router import Router diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 898b296e83e..175432d1dc1 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -10,7 +10,6 @@ legacy internal names with `general_settings.use_team_public_model_name: false`. from __future__ import annotations -import re from collections.abc import Container, Mapping, Sequence from dataclasses import dataclass from functools import reduce @@ -21,15 +20,13 @@ from typing import TYPE_CHECKING, Final, cast from pydantic import TypeAdapter, ValidationError import litellm +from litellm.constants import CLAUDE_CODE_CLIENT, CLAUDE_CODE_PICKER_PATTERN, GATEWAY_CLIENT_HEADER from litellm.router_utils.common_utils import resolve_model_group_alias if TYPE_CHECKING: from litellm.router import Router from litellm.types.proxy.model_listing import ModelInfoResponse -CLAUDE_CODE_PICKER_PATTERN: Final = re.compile(r"claude|anthropic", re.IGNORECASE) -GATEWAY_CLIENT_HEADER: Final = "x-gateway-client" -CLAUDE_CODE_CLIENT: Final = "claude-code" _CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" _ONE_MILLION_SUFFIX: Final = "[1m]" _ONE_MILLION_TOKENS: Final = 1_000_000 diff --git a/litellm/proxy/common_utils/slow_gc_logger.py b/litellm/proxy/common_utils/slow_gc_logger.py new file mode 100644 index 00000000000..63400341426 --- /dev/null +++ b/litellm/proxy/common_utils/slow_gc_logger.py @@ -0,0 +1,35 @@ +"""Garbage collection passes long enough to count as event-loop holds, logged next to asyncio's own +`Executing took X seconds` lines so a reader can tell a collector pause from code that held the loop +(asyncio charges the pause to whichever task was running). Registered only when the loop runs in debug mode.""" + +from __future__ import annotations + +import gc +import time +from collections.abc import Callable, Mapping +from typing import Final, Literal + + +class SlowGcLogger: + def __init__( + self, threshold_seconds: float, log: Callable[[str], None], clock: Callable[[], float] = time.monotonic + ) -> None: + self._threshold_seconds: Final = threshold_seconds + self._log: Final = log + self._clock: Final = clock + self._started = clock() + + def __call__(self, phase: Literal["start", "stop"], info: Mapping[str, int]) -> None: + if phase == "start": + self._started = self._clock() + return + seconds: Final = self._clock() - self._started + if seconds < self._threshold_seconds: + return + self._log( + f"Garbage collection gen={info['generation']} collected={info['collected']} took {seconds:.3f} seconds" + ) + + +def log_slow_gc_passes(threshold_seconds: float, log: Callable[[str], None]) -> None: + gc.callbacks.append(SlowGcLogger(threshold_seconds=threshold_seconds, log=log)) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 78da03150dc..53971ed412c 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -4,6 +4,7 @@ database token (AWS RDS IAM or Microsoft Entra ID) fresh. """ import asyncio +import importlib.resources import os import random import signal @@ -37,6 +38,8 @@ if TYPE_CHECKING: from prisma import Prisma from prisma.types import HttpConfig +_ENVIRONMENT_VARIABLES_REVISION_MIGRATION: Final = "20261015120000_add_environment_variables_revision_trigger" + __all__ = ( "IAMEndpoint", "PrismaManager", @@ -990,6 +993,36 @@ class PrismaManager: return ProxyExtrasDBManager.apply_replica_identity_full_if_requested() + @staticmethod + def _apply_environment_variables_revision_trigger() -> None: + try: + from litellm_proxy_extras.prisma_toolchain import ( + prisma_command_timeout, + run_prisma, + ) + except ImportError: + return + migration: Final = ( + importlib.resources.files("litellm_proxy_extras") + / "migrations" + / _ENVIRONMENT_VARIABLES_REVISION_MIGRATION + / "migration.sql" + ) + if not migration.is_file(): + return + try: + run_prisma( + ["prisma", "db", "execute", "--schema", "schema.prisma", "--file", str(migration)], + timeout=prisma_command_timeout(), + env=os.environ.copy(), + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e: + verbose_proxy_logger.warning( + "Could not create the environment_variables revision trigger after prisma db push, so a removed " + "DB environment variable stays set on running proxies until they restart: %s", + e, + ) + @staticmethod def _raise_if_partitioned_spend_logs() -> None: """`prisma db push` rewrites a doc-partitioned LiteLLM_SpendLogs @@ -1066,6 +1099,7 @@ class PrismaManager: stdout=None, stderr=None, ) + PrismaManager._apply_environment_variables_revision_trigger() PrismaManager._apply_replica_identity_full_if_requested() return True except subprocess.TimeoutExpired as e: diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index c3a6f259f95..2baeabdce7b 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -8,8 +8,13 @@ hook scan such a stream, and re-emit it when the guardrail rewrote the response. from __future__ import annotations import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details +from litellm.llms.anthropic.pass_through.stream_assembly import anthropic_message_delta_fields from litellm.llms.anthropic.pass_through.stream_assembly import ( assemble_anthropic_sse_stream as assemble_anthropic_sse_stream, # noqa: PLC0414 # public re-export ) @@ -50,7 +55,62 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402 ) -def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]: +def _tool_use_inputs(content: Sequence[object]) -> Mapping[str, object]: + blocks: Final = (block for raw in content if (block := as_str_mapping(raw)) is not None) + return MappingProxyType( + { + tool_use_id: block.get("input") + for block in blocks + if block.get("type") == "tool_use" and isinstance(tool_use_id := block.get("id"), str) + } + ) + + +def _unchanged_tool_use_ids(source_content: Sequence[object], delivered_content: Sequence[object]) -> frozenset[str]: + source_inputs: Final = _tool_use_inputs(source_content) + return frozenset( + tool_use_id + for tool_use_id, tool_input in _tool_use_inputs(delivered_content).items() + if tool_use_id in source_inputs and source_inputs[tool_use_id] == tool_input + ) + + +def _rewritten_message_delta_fields( + source_chunks: Sequence[object], unchanged_tool_use_ids: frozenset[str], keep_explanations: bool +) -> Mapping[str, object]: + source_fields: Final = anthropic_message_delta_fields(source_chunks) + stop_details: Final = source_fields.get("stop_details") + safeguard_results: Final = source_fields.get("safeguard_results") + return MappingProxyType( + { + **( + {"stop_details": redacted_stop_details(stop_details, keep_explanations)} + if stop_details is not None + else {} + ), + **( + { + "safeguard_results": filtered_safeguard_results( + safeguard_results, unchanged_tool_use_ids, keep_explanations + ) + } + if safeguard_results is not None + else {} + ), + } + ) + + +def anthropic_sse_chunks_from_response( + assembled: ModelResponse, source_chunks: Sequence[object] = (), *, keep_explanations: bool = False +) -> tuple[bytes, ...]: + """Re-emit a rewritten response as the SSE events of the ``source_chunks`` stream it replaces. + + ``message_delta`` carries the source stream's ``stop_details`` and ``safeguard_results`` and + nothing the OpenAI round trip made up, with a verdict kept only for a tool call delivered with + the input the classifier judged. Their free-text ``explanation`` was written about the text the + guardrail rewrote, so it is dropped unless the caller rewrites no text (``keep_explanations``). + """ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, ) @@ -58,7 +118,22 @@ def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, FakeAnthropicMessagesStreamIterator, ) - anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( - response=assembled + adapter: Final = LiteLLMAnthropicMessagesAdapter() + anthropic_response: Final = adapter.translate_openai_response_to_anthropic(response=assembled) + source_response: Final = assemble_anthropic_sse_stream(source_chunks) + source_content: Final = ( + adapter.translate_openai_response_to_anthropic(response=source_response).get("content") + if source_response is not None + else None + ) + unchanged_tool_use_ids: Final = _unchanged_tool_use_ids( + source_content or (), anthropic_response.get("content") or () + ) + return tuple( + FakeAnthropicMessagesStreamIterator( + response=anthropic_response, + message_delta_fields=_rewritten_message_delta_fields( + source_chunks, unchanged_tool_use_ids, keep_explanations + ), + ).chunks ) - return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index b92020c89be..0672ccf6a66 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -3058,7 +3058,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################################### if raw_sse: for sse_chunk in ( - anthropic_sse_chunks_from_response(assembled_model_response) + anthropic_sse_chunks_from_response(assembled_model_response, all_chunks) if model_response_text(assembled_model_response) != pre_guardrail_text else all_chunks ): 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 e590e5c2e43..b9871b1e276 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1263,7 +1263,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Return sanitized stream if surface is _StreamSurface.ANTHROPIC_MESSAGES: - for sse_chunk in anthropic_sse_chunks_from_response(assembled_response): + for sse_chunk in anthropic_sse_chunks_from_response(assembled_response, all_chunks): yield sse_chunk return mock_response: Final = MockResponseIterator(model_response=assembled_response) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 360752d201a..718401b6e68 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -4,12 +4,13 @@ # # +-------------------------------------------------------------+ +import asyncio import enum import json import os from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeAlias, cast from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger @@ -56,6 +57,11 @@ class _Action(str, enum.Enum): class NomaV2Guardrail(CustomGuardrail): + reads_stream_marks: ClassVar[bool] = True + """Noma keeps one record per streamed reply when every scan carries ``stream`` and only the last one carries the + finished ``response``, as chat completions sends them, so provider streams relayed on the passthrough route mark + this guardrail's scans that way.""" + def __init__( self, api_key: str | None = None, @@ -361,6 +367,10 @@ class NomaV2Guardrail(CustomGuardrail): response_json if isinstance(response_json, dict) else getattr(e, "detail", blocked_detail) ) raise + except asyncio.CancelledError: + # The caller left before Noma answered, so there is no verdict to record. + guardrail_status = "not_run" # rebind-ok: the status the finally block records + raise except Exception as e: guardrail_status = "guardrail_failed_to_respond" guardrail_json_response = str(e) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 01a9259203a..cadc9ca689b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -16,6 +16,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaita from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime +from itertools import accumulate, count, pairwise from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast import aiohttp @@ -182,6 +183,59 @@ async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGener yield chunk +@dataclass(frozen=True, slots=True) +class _Detection: + entity_type: str + start: int + end: int + score: float + + +def _strictly_contains(outer: _Detection, inner: _Detection) -> bool: + return ( + outer.start <= inner.start and inner.end <= outer.end and (outer.start, outer.end) != (inner.start, inner.end) + ) + + +def _collapse_run(run: Sequence[_Detection]) -> _Detection: + label: Final = max( + (detection for detection in run if not any(_strictly_contains(other, detection) for other in run)), + key=lambda detection: (detection.score, detection.end - detection.start, detection.entity_type), + ) + return _Detection(label.entity_type, run[0].start, max(detection.end for detection in run), label.score) + + +def _to_detection(item: PresidioAnalyzeResponseItem) -> _Detection: + return _Detection( + str(item.get("entity_type")), + int(item.get("start") or 0), + int(item.get("end") or 0), + float(item.get("score") or 0.0), + ) + + +def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponseItem]) -> tuple[_Detection, ...]: + """ + Collapse each run of overlapping detections into one detection covering the whole run. + + The run takes its entity type from its highest-scoring detection that no other + detection in the run strictly contains, the same containment rule Presidio's + anonymizer applies, so a phone number with a date-shaped fragment inside it stays + a PHONE_NUMBER. Runs come back sorted by start and never overlap. + """ + detections: Final = sorted( + (detection for detection in map(_to_detection, analyze_results) if detection.start < detection.end), + key=lambda detection: detection.start, + ) + run_ends: Final = tuple(accumulate((detection.end for detection in detections), max)) + boundaries: Final = ( + 0, + *(index for index in range(1, len(detections)) if detections[index].start >= run_ends[index - 1]), + len(detections), + ) + return tuple(_collapse_run(detections[first:last]) for first, last in pairwise(boundaries) if first < last) + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -822,31 +876,25 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data["metadata"] = {} if "pii_tokens" not in request_data["metadata"]: request_data["metadata"]["pii_tokens"] = {} - pii_tokens: Final = request_data["metadata"]["pii_tokens"] + pii_tokens: Final = cast("dict[str, str]", request_data["metadata"]["pii_tokens"]) # cast-ok: this hook's map - # Assign sequence numbers in forward (left-to-right) order so - # that is the first entity in the text, etc. - sorted_forward: Final = sorted(analyze_results, key=lambda x: x["start"]) - seq_map: Final = {} - for idx, ar in enumerate(sorted_forward, start=1): - seq_map[(ar["start"], ar["end"])] = idx + runs: Final = _resolve_overlapping_spans(analyze_results) + # Every message and content part numbers its runs from 1 into this one map, so a + # token another part already holds for a different value takes a spare number. + # No token in the map is numbered above its size, so spares start above both. + spare_numbers: Final = count(max(len(pii_tokens), len(runs)) + 1) - # Apply replacements in reverse order by start position so - # that replacing later spans first does not shift earlier - # coordinates in the original text. - for ar in reversed(sorted_forward): - start = ar["start"] - end = ar["end"] - entity_type = ar["entity_type"] - replacement = f"<{entity_type}>" - seq = seq_map[(start, end)] - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" - pii_tokens[replacement] = text[start:end] - new_text = new_text[:start] + replacement + new_text[end:] - masked_entity_count[entity_type] = masked_entity_count.get(entity_type, 0) + 1 + for seq, run in reversed(tuple(enumerate(runs, start=1))): + original = text[run.start : run.end] + numbered = f"<{run.entity_type}_{seq}>" + replacement = ( + numbered + if pii_tokens.get(numbered, original) == original + else f"<{run.entity_type}_{next(spare_numbers)}>" + ) + pii_tokens[replacement] = original + new_text = new_text[: run.start] + replacement + new_text[run.end :] + masked_entity_count[run.entity_type] = masked_entity_count.get(run.entity_type, 0) + 1 return new_text async def anonymize_text( @@ -1532,7 +1580,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): await self._process_response_for_pii(response=assembled, request_data=request_data, mode="mask") if model_response_text(assembled) == original_text: return chunks - return anthropic_sse_chunks_from_response(assembled) + return anthropic_sse_chunks_from_response(assembled, chunks) @staticmethod def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes: diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index ebd50c8cd39..fa92a0c90e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -971,7 +971,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return self._modify_response_with_permission_errors(anthropic_response, anthropic_denials) - for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response): + for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response, all_chunks, keep_explanations=True): yield sse_chunk def _check_assembled_stream( 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 a7f813d953e..3ad9d1a2c42 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -82,6 +82,9 @@ class _EndpointTranslation(Protocol): @property def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... + @property + def needs_ended_stream_scan(self) -> "Callable[[CustomGuardrail, Sequence[object]], bool]": ... + @property def released_stream_as_ended(self) -> "Callable[[Sequence[object]], tuple[object, ...]]": ... @@ -153,6 +156,21 @@ def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "Stre return scan_key == last_scan_key or scan_key.has_nothing_to_scan +def _owes_ended_scan( + endpoint_translation: _EndpointTranslation, + guardrail: CustomGuardrail, + responses_so_far: Sequence[object], + last_scan_key: "StreamingScanKey | None", +) -> bool: + """Whether the end-of-stream round still runs after a sampled round covered every text before the stream ended, + because the translation hands ``guardrail`` marks only an ended round carries.""" + return ( + last_scan_key is not None + and not last_scan_key.stream_ended + and endpoint_translation.needs_ended_stream_scan(guardrail, responses_so_far) + ) + + 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.""" @@ -1233,6 +1251,13 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is None: call_type = infer_call_type(call_type=None, completion_response=item) + # Provider passthrough routes (e.g. /bedrock/.../converse-stream) map to no call type, and their + # frames carry none; the logging object names the passthrough, as on the non-streaming path. + if call_type is None: + logging_call_type = getattr(typed_request_data.get("litellm_logging_obj"), "call_type", None) + if logging_call_type == CallTypes.allm_passthrough_route.value: + call_type = logging_call_type + # If call type not supported, just pass through all chunks if call_type is None or CallTypes(call_type) not in mappings: yield item @@ -1383,7 +1408,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) verdict_settled = True - if _is_redundant_scan(end_scan_key, last_scan_key): + if _is_redundant_scan(end_scan_key, last_scan_key) and not _owes_ended_scan( + endpoint_translation, guardrail_to_apply, responses_so_far, 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, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 592c732bfd3..54f4ca25974 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -322,9 +322,9 @@ def _credentials_table(prisma_client: PrismaClient) -> TableActions[CredentialIt ) -def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions: +def _config_table(prisma_client: PrismaClient, *, use_writer: bool = False) -> _ConfigTableActions: return cast( # cast-ok: ConfigRepository.table hides the write actions this module needs on that same object - "_ConfigTableActions", ConfigRepository(prisma_client).table + "_ConfigTableActions", ConfigRepository(prisma_client, use_writer=use_writer).table ) @@ -5566,7 +5566,7 @@ async def rotate_master_key( await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") # 3. process config table try: - config = await _config_table(prisma_client).find_many() + config = await _config_table(prisma_client, use_writer=True).find_many() except Exception: config = None @@ -5589,7 +5589,7 @@ async def rotate_master_key( if encrypted_env_vars: await _config_table(prisma_client).update( where={"param_name": "environment_variables"}, - data={"param_value": prisma.Json(encrypted_env_vars)}, + data={"param_value": prisma.Json({**environment_variables_dict, **encrypted_env_vars})}, ) try: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d014ff6cf07..a014b9cf8a5 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2517,7 +2517,9 @@ if MCP_AVAILABLE: resolved_redirect_uri: Final = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri caller_client_id: Final = sealed_code.client_id if sealed_code else client_id caller_client_secret: Final = sealed_code.client_secret if sealed_code else client_secret - resolved_client_id: Final = mcp_server.client_id or caller_client_id or "" + resolved_client_id: Final = ( + caller_client_id if sealed_code else (mcp_server.client_id or caller_client_id or "") + ) if not resolved_client_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -2541,6 +2543,7 @@ if MCP_AVAILABLE: refresh_token=refresh_token, scope=scope, client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None, + sealed_client=sealed_code is not None, ) @router.post( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 464cfee270a..3aef5724bdc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1087,6 +1087,16 @@ def _guardrail_refusal_status(error: Exception) -> int | None: return None +def _aws_error_type(error: Exception, handled: ProxyException) -> str | None: + """The AWS error type to answer the handled error with: ValidationException for a guardrail block, else the type + the error carried, as LiteLLM's own Bedrock request checks set it. None when the shared handler changed the + guardrail block's status.""" + refusal_status: Final = _guardrail_refusal_status(error) + if refusal_status is not None: + return "ValidationException" if str(refusal_status) == handled.code else None + return next((value for name, value in handled.headers.items() if name.lower() == AWS_ERROR_TYPE_HEADER), None) + + def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None: """Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error type.""" @@ -1111,9 +1121,9 @@ async def bedrock_passthrough_error_response( user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLoggingType, ) -> Response: - """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error - Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the - shared handler raises it.""" + """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block, an error + Bedrock returned, or an error carrying an AWS error type the way Bedrock does, so an AWS SDK reads its code and + message. Anything else is raised as the shared handler raises it.""" try: await processor.handle_llm_api_exception( e=error, @@ -1126,9 +1136,9 @@ async def bedrock_passthrough_error_response( raise return response except ProxyException as handled: - refusal_status: Final = _guardrail_refusal_status(error) - if refusal_status is not None and str(refusal_status) == handled.code: - raise AwsErrorTypeProxyException(handled, "ValidationException") from handled + aws_error_type: Final = _aws_error_type(error, handled) + if aws_error_type is not None: + raise AwsErrorTypeProxyException(handled, aws_error_type) from handled raise diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 3c5f55d87a8..7f68b72a522 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -713,6 +713,8 @@ class ProxyInitializationHelpers: """Helper function to determine the event loop type based on platform""" if sys.platform in ("win32", "cygwin", "cli"): return None # Let uvicorn choose the default loop on Windows + if os.environ.get("LITELLM_EVENT_LOOP") == "asyncio": + return "asyncio" return "uvloop" _get_loop_type = get_loop_type diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3de3dccc532..1e16eb626fe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,9 +1,11 @@ import asyncio import copy import enum +import gc import importlib import inspect import io +import itertools import os import random import re @@ -75,6 +77,7 @@ from litellm.constants import ( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, + INTERNAL_CALLBACK_NAMES, LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, LITELLM_SETTINGS_SAFE_DB_OVERRIDES, @@ -510,6 +513,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import ( parse_stagger_settings, stagger_trigger, ) +from litellm.proxy.common_utils.slow_gc_logger import log_slow_gc_passes from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, @@ -551,6 +555,7 @@ from litellm.proxy.config_resolvers.settings_rules import ( from litellm.proxy.config_resolvers.settings_rules import ( JsonValue as SettingsJsonValue, ) +from litellm.proxy.config_resolvers.sso import SSO_FIELD_ENV_VARS from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.create_views import SupportsRawQueries @@ -1379,6 +1384,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState init_verbose_loggers() + running_loop: Final = asyncio.get_running_loop() + if isinstance(running_loop, asyncio.BaseEventLoop) and running_loop.get_debug(): + log_slow_gc_passes(threshold_seconds=running_loop.slow_callback_duration, log=verbose_proxy_logger.warning) + prometheus_multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR") if prometheus_multiproc_dir: mark_dead_workers(prometheus_multiproc_dir) @@ -1738,6 +1747,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState async with AsyncExitStack() as admin_mcp_stack: try: await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app)) + gc.freeze() + verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count()) yield state finally: if model_info_scheduler is not None and model_info_scheduler.running: @@ -2225,9 +2236,9 @@ class _InvitationAcceptTxManager(Protocol): _ProxyModelRow: TypeAlias = "prisma_models.LiteLLM_ProxyModelTable" -def _config_param_table(client: PrismaClient | None) -> TableActions[_ConfigParamRow]: +def _config_param_table(client: PrismaClient | None, *, use_writer: bool = False) -> TableActions[_ConfigParamRow]: return cast( # cast-ok: this is prisma's LiteLLM_Config actions object, which parses its Json column to a mapping - "TableActions[_ConfigParamRow]", ConfigRepository(client).table + "TableActions[_ConfigParamRow]", ConfigRepository(client, use_writer=use_writer).table ) @@ -5565,6 +5576,23 @@ def _as_settings_mapping(value: object) -> Mapping[str, SettingsJsonValue]: return _SETTINGS_MAPPING.validate_python(value) +_SSO_ENVIRONMENT_VARIABLES: Final = frozenset(SSO_FIELD_ENV_VARS.values()) +_ENVIRONMENT_VALUES: Final = TypeAdapter(Mapping[str, str]) + + +def _config_row_revision(row: object) -> int | None: + revision: Final[object] = getattr(row, "reload_revision", None) + return revision if isinstance(revision, int) and not isinstance(revision, bool) else None + + +def _restore_environment(keys: frozenset[str], values: Mapping[str, str]) -> None: + for key in keys: + if (value := values.get(key)) is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + def _get_field_default(field_info: FieldInfo) -> JsonValue: if field_info.default is PydanticUndefined: return None @@ -5726,6 +5754,12 @@ class ProxyConfig: self.router_settings: Final[SettingsStore] = SettingsStore("router_settings") self.litellm_settings: Final[SettingsStore] = SettingsStore("litellm_settings") self.environment_variables: Final[SettingsStore] = SettingsStore("environment_variables") + self._environment_before_db: Mapping[str, str] | None = None + self._db_environment_revision: int | None = None + self._environment_read_numbers: Final = itertools.count(1) + self._latest_applied_environment_read: int = 0 + self._db_environment_keys: frozenset[str] = frozenset() + self._sso_environment_removals: frozenset[str] = frozenset() self._warned_shadowed_keys: frozenset[tuple[Section, str]] = frozenset() self._db_config_callback_entries: Mapping[tuple[str, str], _DbCallbackRegistration] = MappingProxyType({}) self._settings_stores: Final[Mapping[Section, SettingsStore]] = MappingProxyType( @@ -6030,7 +6064,7 @@ class ProxyConfig: if prisma_client is None or not (general_settings.get("store_model_in_db", False) is True or store_model_in_db): return - row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( + row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client, use_writer=True).find_first( where={"param_name": "environment_variables"} ) existing: Final[dict] = dict(row.param_value) if row is not None and row.param_value is not None else {} @@ -8364,6 +8398,8 @@ class ProxyConfig: verbose_proxy_logger.info("'store_model_in_db' is not True, skipping db updates") return config + environment_before_db: Final = self._captured_environment_before_db() + environment_read: Final = next(self._environment_read_numbers) sections: Final = tuple(self._settings_stores) responses: Final = await asyncio.gather(*(get_config_param(prisma_client, section) for section in sections)) for section, response in zip(sections, responses): @@ -8377,6 +8413,10 @@ class ProxyConfig: ) if section == "litellm_settings": self._apply_litellm_settings_db_values(self._prepared_db_settings_values(section, param_value)) + elif section == "environment_variables": + await self._apply_db_environment_variables_row( + prisma_client, response, environment_before_db, environment_read + ) else: self._settings_stores[section].apply_db_row( section, @@ -8386,6 +8426,49 @@ class ProxyConfig: self._warn_about_shadowed_db_settings() return self._config_with_resolved_settings(config) + async def _apply_db_environment_variables_row( + self, prisma_client: PrismaClient, row: object, environment_before_db: Mapping[str, str], read: int + ) -> None: + current: Final = await self._environment_row_not_behind_the_applied_one(prisma_client, row) + if current is None or (param_value := getattr(current, "param_value", None)) is None: + return + self._apply_db_environment_variables(param_value, _config_row_revision(current), environment_before_db, read) + + async def _environment_row_not_behind_the_applied_one( + self, prisma_client: PrismaClient, row: object + ) -> object | None: + if not self._is_behind_applied_environment_revision(row): + return row + applied_revision: Final = self._db_environment_revision + try: + writer_row: Final = await call_with_db_reconnect_retry( + prisma_client, + lambda: ConfigRepository(prisma_client, use_writer=True).table.find_unique( + where={"param_name": "environment_variables"} + ), + reason="environment_variables_writer_read_failure", + ) + except Exception as e: + if not ( + PrismaDBExceptionHandler.is_database_transport_error(e) + or PrismaDBExceptionHandler.is_database_connection_error(e) + ): + raise + verbose_proxy_logger.warning( + "Could not read environment_variables from the writer database, keeping the current environment " + "variables until the next config sync: %s", + e, + ) + return None + return writer_row if self._db_environment_revision == applied_revision else None + + def _is_behind_applied_environment_revision(self, row: object) -> bool: + applied_revision: Final = self._db_environment_revision + if applied_revision is None or row is None: + return False + revision: Final = _config_row_revision(row) + return revision is None or revision < applied_revision + def _warn_about_shadowed_db_settings(self) -> None: shadowed: Final[frozenset[tuple[Section, str]]] = frozenset( (section, key) for section, store in self._settings_stores.items() for key in store.shadowed_db_keys() @@ -8397,21 +8480,47 @@ class ProxyConfig: self._warned_shadowed_keys = shadowed def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]: - if section == "environment_variables": - decrypted: Final = self.decrypt_and_set_db_env_variables( - dict(_as_settings_mapping(value)), return_original_value=True - ) - normalized: Final = { - **decrypted, - **{key.upper(): decrypted_value for key, decrypted_value in decrypted.items()}, - } - for key, decrypted_value in normalized.items(): - os.environ[key] = decrypted_value - return _as_settings_mapping(normalized) - scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=value) return _as_settings_mapping(scrubbed) + def _apply_db_environment_variables( + self, value: object, revision: int | None, environment_before_db: Mapping[str, str], read: int + ) -> None: + if read < self._latest_applied_environment_read: + return + self._latest_applied_environment_read = read + applied_revision: Final = self._db_environment_revision + row: Final = _as_settings_mapping(value) + decrypted: Final = _ENVIRONMENT_VALUES.validate_python( + self.decrypt_and_set_db_env_variables(dict(row), return_original_value=True) + ) + applied: Final = {**decrypted, **{key.upper(): decrypted_value for key, decrypted_value in decrypted.items()}} + os.environ.update(applied) + if revision is None or revision == applied_revision: + self._db_environment_keys = self._db_environment_keys | frozenset(applied) + else: + stored: Final = frozenset(key for key, stored_value in row.items() if stored_value is not None) + still_stored: Final = stored | frozenset(key.upper() for key in stored) + removed: Final = self._db_environment_keys - still_stored + _restore_environment(removed - _SSO_ENVIRONMENT_VARIABLES, environment_before_db) + self._sso_environment_removals = ( + self._sso_environment_removals | (removed & _SSO_ENVIRONMENT_VARIABLES) + ) - still_stored + self._db_environment_keys = (self._db_environment_keys & still_stored) | frozenset(applied) + self._db_environment_revision = revision + self.environment_variables.apply_db_row("environment_variables", _as_settings_mapping(applied)) + + def _restore_sso_environment_removals(self, sso_table_keys: frozenset[str]) -> None: + _restore_environment( + self._sso_environment_removals - sso_table_keys, self._environment_before_db or MappingProxyType({}) + ) + self._sso_environment_removals = frozenset() + + def _captured_environment_before_db(self) -> Mapping[str, str]: + if self._environment_before_db is None: + self._environment_before_db = MappingProxyType(dict(os.environ)) + return self._environment_before_db + def _apply_litellm_settings_db_values(self, db_values: Mapping[str, SettingsJsonValue]) -> None: self.litellm_settings.apply_db_row("litellm_settings", db_values) for key in LITELLM_SETTINGS_SAFE_DB_OVERRIDES: @@ -8645,6 +8754,8 @@ class ProxyConfig: if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) + else: + self._restore_sso_environment_removals(frozenset()) if self._should_load_db_object(object_type="cache_settings"): from litellm.proxy.management_endpoints.cache_settings_endpoints import ( CacheSettingsManager, @@ -8821,6 +8932,11 @@ class ProxyConfig: sso_settings.sso_settings.pop("ui_access_mode", None) uppercase_sso_settings: Final = {key.upper(): value for key, value in sso_settings.sso_settings.items()} self.decrypt_and_set_db_env_variables(environment_variables=uppercase_sso_settings) + self._restore_sso_environment_removals( + frozenset() + if sso_settings is None + else frozenset(key.upper() for key, value in sso_settings.sso_settings.items() if value) + ) except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - %s", e @@ -18783,7 +18899,7 @@ async def update_config( proxy_config.reject_config_owned_writes(section_name="router_settings", changed_keys=router_settings_updates) async def _read_section(param_name: str) -> dict: - row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( + row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client, use_writer=True).find_first( where={"param_name": param_name} ) if row is None or row.param_value is None: @@ -19922,8 +20038,7 @@ def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) -> if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE ) configured_modules: Final = frozenset(name.rsplit(".", 1)[0] for name in configured_callback_names if "." in name) - internal_callback_names: Final = frozenset({"cache", "vector_store_pre_call_hook"}) - return internal_callback_names | frozenset( + return INTERNAL_CALLBACK_NAMES | frozenset( callback_name for callback_name, callback in litellm.logging_callback_manager.get_callback_objects() if isinstance(callback, CustomGuardrail) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 4a6a88f9a86..cb9e193fa48 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2417,6 +2417,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "Nace", + "provider_display_name": "Nace", + "litellm_provider": "nace", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://console.nace.ai", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "nace/drex-latest" + }, { "provider": "NEBIUS", "provider_display_name": "Nebius", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 778d40c2d07..15888d2601a 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,6 +1,7 @@ import asyncio import contextlib import json +import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum @@ -28,10 +29,11 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.websockets import WebSocket, WebSocketDisconnect from typing_extensions import ReadOnly, TypedDict +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.tool_fields import first_validation_error @@ -58,12 +60,18 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) +from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id +from litellm.proxy.litellm_pre_call_utils import ( + _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list +) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import handle_exception_on_proxy +from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1813,6 +1821,143 @@ async def _enforce_responses_ws_first_frame_model_auth( ) +_WS_FRAME: Final = TypeAdapter(dict[str, object]) +_GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]]) +_PII_TOKENS: Final = TypeAdapter(dict[str, str]) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" +_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") +_CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") +# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it +_CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") + + +class _PreCallGuardrailWalk(Protocol): + async def pre_call_hook( + self, + *, + user_api_key_dict: UserAPIKeyAuth, + data: dict[str, object], # mutable-ok: ProxyLogging.pre_call_hook hands the guardrails a dict they edit + call_type: Literal["aresponses"], + guardrails_only: bool, + ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + + +def _responses_ws_frame_guard( + admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: _PreCallGuardrailWalk +) -> ResponsesWsFrameGuard: + """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every + ``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) + admitted_bucket: Final = admitted.get(metadata_key) + bucket: Final = { + key: value + for key, value in ( + _WS_FRAME.validate_python(admitted_bucket) if isinstance(admitted_bucket, dict) else {} + ).items() + if key != _GUARDRAIL_INFORMATION_KEY + } + model: Final = admitted.get("model") + + # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation + rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + + async def guard(frame: str) -> GuardedResponsesWsFrame: + frame_obj: Final = _WS_FRAME.validate_json(frame) + nested: Final = frame_obj.get("response") + if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): + raise ValueError( + "response.create carries input, instructions or tools both at the top level and under response" + ) + body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj + frame_data: Final = { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied + metadata_key: {**bucket}, + } + try: + guarded_data: Final = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True + ) + except Exception as e: + await _bill_rejected_frame(rejected_frame_auth, e, frame_data) + raise + if guarded_data.get("model") != model: + rerouted: Final = ValueError( + f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" + f" is bound to {model!r}, send this turn over HTTP /v1/responses" + ) + await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) + raise rerouted + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} + kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} + rewritten: Final = ( + {**frame_obj, "response": {**kept, **guarded_fields}} + if isinstance(nested, dict) + else {**kept, **guarded_fields} + ) + return GuardedResponsesWsFrame( + frame=json.dumps(rewritten), + guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)), + pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + ) + + return guard + + +def _without_untrusted_metadata(key: str, value: object) -> object: + """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" + if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): + return value + return { + name: item + for name, item in _WS_FRAME.validate_python(value).items() + if not name.startswith("user_api_key_") + and name not in _UNTRUSTED_METADATA_CONTROL_FIELDS + and name not in _CLIENT_PRICING_METADATA_FIELDS + and name not in ("headers", "pii_tokens", *_CLIENT_GUARDRAIL_SELECTORS) + } + + +async def _bill_rejected_frame( + user_api_key_dict: UserAPIKeyAuth, + exception: Exception, + frame_data: dict[str, object], # mutable-ok: the spend logger stamps the failure row on the request dict +) -> None: + """Only the spend logger sees a rejected frame; the other failure hooks (the rate limiter's slot release, + alerting) belong to the live connection, which stays open. Billing never replaces the guardrail's own error.""" + for callback in litellm.callbacks: + if isinstance(callback, ProxyDBLogger): + try: + await callback.async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict + request_data=frame_data, original_exception=exception, user_api_key_dict=user_api_key_dict + ) + except Exception: # noqa: BLE001 # billing never replaces the guardrail's own error + verbose_proxy_logger.exception("[Non-Blocking] Responses WebSocket: billing a rejected frame failed") + + +def _recorded_pii_tokens(metadata: object) -> Mapping[str, str]: + """Presidio stores the tokens it masked a frame with in ``metadata.pii_tokens``; the socket unmasks with them.""" + recorded: Final = _WS_FRAME.validate_python(metadata).get("pii_tokens") if isinstance(metadata, dict) else None + return MappingProxyType(_PII_TOKENS.validate_python(recorded)) if isinstance(recorded, dict) else _NO_PII_TOKENS + + +_NO_PII_TOKENS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _recorded_guardrail_information(bucket: object) -> tuple[Mapping[str, object], ...]: + recorded: Final = ( + _WS_FRAME.validate_python(bucket).get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + ) + return tuple(_GUARDRAIL_INFORMATION.validate_python(recorded)) if recorded else () + + async def _responses_websocket_session( websocket: WebSocket, model: str | None, @@ -1888,6 +2033,7 @@ async def _responses_websocket_session( user_api_base=user_api_base, model=resolved_model, route_type="_aresponses_websocket", + skip_guardrails=True, ) except Exception as e: verbose_proxy_logger.exception("Responses WebSocket pre-call error") @@ -1909,7 +2055,10 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_frame_guard=_responses_ws_frame_guard(data, user_api_key_dict, proxy_logging_obj), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 4ece88f757d..d51e1167888 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -1326,7 +1326,7 @@ async def update_sso_settings( # Remove SSO-related env vars from config.environment_variables try: - env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique( + env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client, use_writer=True)).find_unique( where={"param_name": "environment_variables"} ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 36ace3c1697..29749afb429 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4167,6 +4167,49 @@ class ProxyLogging: await ProxyLogging._record_served_stream_output(request_data, served_chunks) ProxyLogging.fire_deferred_stream_logging(request_data) + async def async_post_call_event_stream_iterator_hook( + self, + response: AsyncIterable[object], + user_api_key_dict: UserAPIKeyAuth, + request_data: dict[str, object], # mutable-ok: same request-payload shape the hooks mutate + guardrails: Sequence[tuple[CustomGuardrail, bool]], + ) -> AsyncGenerator[object, None]: + """ + Run post_call guardrails over a provider event stream the proxy relays as it arrives, such as a Bedrock + Converse stream on the passthrough route. Each guardrail scans the stream through the unified guardrail, as + on /v1/messages; its flag says whether it holds the stream until moderated by default. + """ + from litellm.proxy.proxy_server import llm_router + + merged_request_data: Final[dict[str, object]] = check_and_merge_model_level_guardrails( # pyright: ignore[reportUnknownVariableType] # returns the untyped request dict # mutable-ok: same request-payload shape the hooks mutate + data=request_data, llm_router=llm_router + ) + current_response: AsyncIterable[object] = response # rebind-ok: each guardrail wraps the previous layer + guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect + for guardrail, holds_until_moderated in guardrails: + if guarded_layers: + current_response = ProxyLogging._released_copies(current_response) + guarded_layers.append(current_response) + current_response = self._wrap_streaming_iterator_with_enrichment( + guardrail, + current_response, + partial( + unified_guardrail.async_post_call_streaming_iterator_hook, # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + user_api_key_dict=user_api_key_dict, + request_data=merged_request_data, + guardrail_to_apply=guardrail, + buffer_until_moderated_default=holds_until_moderated, + ), + request_data=merged_request_data, + ) + guarded_layers.append(current_response) + try: + async for item in current_response: + yield item + except (GeneratorExit, asyncio.CancelledError): + await ProxyLogging._close_guarded_layers(guarded_layers) + raise + async def _pipeline_gated_stream( self, response: "AsyncGenerator[object, None]", @@ -4351,6 +4394,13 @@ class ProxyLogging: event_type=GuardrailEventHooks.post_call, ) + @staticmethod + async def _released_copies(items: AsyncIterable[object]) -> AsyncGenerator[object, None]: + """Each item a guardrail layer released, copied for the next layer, so a rewrite the earlier layer writes + back once its stream ends never reaches items the next layer is still holding.""" + async for item in items: + yield copy.copy(item) + @staticmethod async def _close_guarded_layers(layers: Sequence[AsyncGenerator[object, None]]) -> None: for layer in reversed(layers): @@ -4559,11 +4609,12 @@ litellm_config_cache: Final[DualCache] = DualCache( class _ConfigRow: """Mimics the Prisma litellm_config row shape for cached entries.""" - __slots__ = ("param_name", "param_value") + __slots__ = ("param_name", "param_value", "reload_revision") - def __init__(self, param_name: str, param_value: object) -> None: + def __init__(self, param_name: str, param_value: object, reload_revision: int | None = None) -> None: self.param_name = param_name self.param_value = param_value + self.reload_revision = reload_revision CONFIG_PARAMS_TARGET: Final = "config_params" @@ -4574,14 +4625,23 @@ def _config_cache_key(param_name: str) -> str: def _pack_config_row(row: Any) -> dict[str, object]: - return {"param_name": row.param_name, "param_value": row.param_value} + return { + "param_name": row.param_name, + "param_value": row.param_value, + "reload_revision": getattr(row, "reload_revision", None), + } def _unpack_config_row(cached: object) -> _ConfigRow | None: if cached is None or cached == _CONFIG_CACHE_MISS: return None if isinstance(cached, dict): - return _ConfigRow(cached["param_name"], cached["param_value"]) + revision: Final = cached.get("reload_revision") + return _ConfigRow( + cached["param_name"], + cached["param_value"], + revision if isinstance(revision, int) and not isinstance(revision, bool) else None, + ) return None diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 4acce4de586..96b4b081c05 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -37,9 +37,10 @@ from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ..litellm_core_utils.realtime_streaming import client_websocket_scope from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client -from ..llms.bedrock.realtime.handler import BedrockRealtime +from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed from ..llms.custom_httpx.http_handler import realtime_ssl_for_url from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import OpenAIRealtime @@ -471,6 +472,7 @@ async def _arealtime( For PROXY use only. """ + raise_if_session_committed(client_websocket_scope(websocket)) headers = cast(dict | None, kwargs.get("headers")) extra_headers: Final = cast(dict | None, kwargs.get("extra_headers")) if headers is None: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index fb2daca43fb..29fddceb84b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -62,6 +62,7 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2599,6 +2600,7 @@ async def _aresponses_websocket( api_key: str | None = None, timeout: float | None = None, turn_admission: Callable[[], Awaitable[None]] | None = None, + responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: """ @@ -2692,5 +2694,6 @@ async def _aresponses_websocket( custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), turn_admission=turn_admission, + frame_guard=responses_ws_frame_guard, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index bcd3fef9415..4c83436218d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -8,12 +8,23 @@ import time import traceback import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from functools import lru_cache from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Protocol, + TypeAlias, + cast, + overload, + runtime_checkable, +) import httpx from openai._streaming import SSEDecoder @@ -29,6 +40,7 @@ from litellm.constants import ( from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -265,7 +277,11 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: def _map_stream_error_to_exception( - error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None + error_obj: object, + model: str, + custom_llm_provider: str, + headers: httpx.Headers | None = None, + logging_obj: LiteLLMLoggingObj | None = None, ) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -278,13 +294,18 @@ def _map_stream_error_to_exception( body=error_body, headers=headers, ) + litellm_params: Final = ( + None + if logging_obj is None + else _typed_gets_litellm_params(logging_obj.model_call_details.get)("litellm_params", {}) + ) try: return litellm.exception_type( model=model, custom_llm_provider=custom_llm_provider, original_exception=provider_exception, completion_kwargs={}, - extra_kwargs={}, + extra_kwargs=with_deployment_api_base(litellm_params), ) except Exception as mapped_exception: return mapped_exception @@ -677,6 +698,7 @@ class BaseResponsesAPIStreamingIterator: self.model or "", self.custom_llm_provider or "", headers=provider_request_id_headers(self._raw_response_headers), + logging_obj=self.logging_obj, ) def _maybe_raise_for_error_event(self, result: object) -> None: @@ -1959,6 +1981,73 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) _PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +_PII_DELTA_EVENT_TYPES: Final = frozenset( + { + "response.output_text.delta", + "response.reasoning_summary_text.delta", + "response.refusal.delta", + "response.function_call_arguments.delta", + } +) + + +def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], unmask_pii_text: _UnmasksPiiText) -> str: + """Put the original PII back into one client-bound event: ``response.completed`` output text, or a delta.""" + try: + evt_obj: Final = _load_json_object(response_str) + except (json.JSONDecodeError, TypeError): + return response_str + event_type: Final = evt_obj.get("type") + + if event_type == "response.completed": + modified = False + response_obj: Final = evt_obj.get("response") + if not _is_json_object(response_obj): + return response_str + output_items: Final = response_obj.get("output") + for output_item in output_items if _is_json_array(output_items) else []: + if not _is_json_object(output_item): + continue + content = output_item.get("content") + if not _is_json_array(content): + continue + for content_block in content: + if not _is_json_object(content_block): + continue + text = content_block.get("text") + if isinstance(text, str): + unmasked = unmask_pii_text(text, pii_tokens) + if unmasked != text: + content_block["text"] = unmasked + modified = True + return json.dumps(evt_obj) if modified else response_str + + if event_type in _PII_DELTA_EVENT_TYPES: + delta: Final = evt_obj.get("delta") + if isinstance(delta, str): + unmasked = unmask_pii_text(delta, pii_tokens) + if unmasked != delta: + evt_obj["delta"] = unmasked + return json.dumps(evt_obj) + + return response_str + + +@dataclass(frozen=True, slots=True) +class GuardedResponsesWsFrame: + """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails + may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the + Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply.""" + + frame: str + guardrail_information: tuple[Mapping[str, object], ...] = () + pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + + +ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES @@ -1991,6 +2080,30 @@ def _user_texts_in_input_item(item: object) -> tuple[str, ...]: ) +def _set_turn_guardrail_information( + turn_logging_obj: LiteLLMLoggingObj, guardrail_information: tuple[Mapping[str, object], ...] +) -> None: + """A turn bills the guardrails that ran on its own ``response.create`` frame, never the records the + connection-level logging object carries from an earlier frame.""" + litellm_params: Final = _json_object_or_none( + turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None: + return + for key in _PER_TURN_METADATA_KEYS: + _replace_guardrail_information(litellm_params.get(key), guardrail_information) + + +def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None: + if not _is_json_object(bucket): + return + bucket.pop(_GUARDRAIL_INFORMATION_KEY, None) + if guardrail_information: + bucket[_GUARDRAIL_INFORMATION_KEY] = [ # rebind-ok: the turn copy's own metadata bucket + dict(entry) for entry in guardrail_information + ] + + def _logged_user_messages(client_frame: object) -> tuple[Mapping[str, object], ...]: frame: Final = _json_object_or_none(client_frame) if frame is None or frame.get("type") != "response.create": @@ -2021,13 +2134,17 @@ class _ResponsesWsTurn: def with_event(self, event: Mapping[str, object]) -> _ResponsesWsTurn: return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) - def failure(self, model: str, custom_llm_provider: str) -> Exception | None: + def failure( + self, model: str, custom_llm_provider: str, connection_logging_obj: LiteLLMLoggingObj + ) -> Exception | None: failed_event: Final = next( (event for event in self.events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None ) if failed_event is None: return None - return _map_stream_error_to_exception(_ws_event_error(failed_event), model, custom_llm_provider) + return _map_stream_error_to_exception( + _ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj + ) _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) @@ -2116,11 +2233,13 @@ class ResponsesWebSocketStreaming: custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, turn_admission: Callable[[], Awaitable[None]] | None = None, + frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict + self.frame_guard: ResponsesWsFrameGuard | None = frame_guard self.request_data: dict[str, object] = request_data or {} litellm_metadata: Final = self.request_data.get("litellm_metadata") self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} @@ -2137,12 +2256,15 @@ class ResponsesWebSocketStreaming: self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission - def _start_turn(self, client_frame: str | None) -> None: + def _start_turn( + self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () + ) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time turn_logging_obj: Final = self.logging_obj.copy_for_new_call( litellm_call_id=str(uuid.uuid4()), start_time=start_time ) self._detach_budget_reservation() + _set_turn_guardrail_information(turn_logging_obj, guardrail_information) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) @@ -2161,10 +2283,46 @@ class ResponsesWebSocketStreaming: for key, value in litellm_params.items() } - def _record_client_frame(self, client_frame: str) -> None: + def _record_client_frame( + self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = () + ) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": - self._start_turn(client_frame) + self._start_turn(client_frame, guardrail_information) + + async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None: + """Run the proxy's pre_call guardrails on one ``response.create`` frame. + + Returns None when a guardrail rejects the frame: the client gets an ``error`` event and the + frame is dropped while the connection stays open for the next turn. + """ + frame: Final = _json_object_or_none(message) + if self.frame_guard is None or frame is None or frame.get("type") != "response.create": + return GuardedResponsesWsFrame(frame=message) + try: + return await self.frame_guard(message) + except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event + verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) + await self._send_error_event("invalid_request_error", str(e)) + return None + + async def _forward_guarded_frame(self, message: str) -> None: + guarded: Final = await self._guard_frame(message) + if guarded is None: + return + if guarded.pii_tokens: + self._remember_pii_tokens(guarded.pii_tokens) + masked: Final = await self._mask_response_create(guarded.frame) + self._record_client_frame(masked, guarded.guardrail_information) + await self.backend_ws.send(masked) + + def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None: + metadata: Final = self.request_data.get("metadata") + known: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None + self.request_data["metadata"] = { + **(metadata if _is_json_object(metadata) else {}), + "pii_tokens": {**(known if _is_str_mapping(known) else {}), **pii_tokens}, + } def _record_backend_event(self, event: str) -> None: event_obj: Final = _json_object_or_none(event) @@ -2188,7 +2346,7 @@ class ResponsesWebSocketStreaming: def _log_turn(self, turn: _ResponsesWsTurn) -> None: events: Final = list(turn.events) - exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "") + exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "", self.logging_obj) if exception is None: asyncio.create_task(turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True)) return @@ -2408,14 +2566,7 @@ class ResponsesWebSocketStreaming: return json.dumps(msg_obj) if modified else message # Delta event types whose ``delta`` field may contain PII tokens. - _DELTA_EVENT_TYPES = frozenset( - { - "response.output_text.delta", - "response.reasoning_summary_text.delta", - "response.refusal.delta", - "response.function_call_arguments.delta", - } - ) + _DELTA_EVENT_TYPES = _PII_DELTA_EVENT_TYPES # Terminal events that carry the full output text or tool-call arguments # already delivered by ``response.completed``. Suppressed when output masking @@ -2455,47 +2606,9 @@ class ResponsesWebSocketStreaming: if not pii_tokens: return response_str - try: - evt_obj: Final = _load_json_object(response_str) - except (json.JSONDecodeError, TypeError): - return response_str - cb: Final = self.guardrail_callbacks[0] unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR) - event_type: Final = evt_obj.get("type") - - if event_type == "response.completed": - modified = False - response_obj: Final = evt_obj.get("response") - if not _is_json_object(response_obj): - return response_str - output_items: Final = response_obj.get("output") - for output_item in output_items if _is_json_array(output_items) else []: - if not _is_json_object(output_item): - continue - content = output_item.get("content") - if not _is_json_array(content): - continue - for content_block in content: - if not _is_json_object(content_block): - continue - text = content_block.get("text") - if isinstance(text, str): - unmasked = unmask_pii_text(text, pii_tokens) - if unmasked != text: - content_block["text"] = unmasked - modified = True - return json.dumps(evt_obj) if modified else response_str - - if event_type in self._DELTA_EVENT_TYPES: - delta: Final = evt_obj.get("delta") - if isinstance(delta, str): - unmasked = unmask_pii_text(delta, pii_tokens) - if unmasked != delta: - evt_obj["delta"] = unmasked - return json.dumps(evt_obj) - - return response_str + return _unmask_responses_event(response_str, pii_tokens, unmask_pii_text) async def _mask_response_completed(self, response_str: str) -> str: """ @@ -2607,21 +2720,25 @@ class ResponsesWebSocketStreaming: except Exception: # noqa: BLE001, S110 # client may already be gone pass + async def _send_error_event(self, error_type: str, message: str) -> None: + try: + await self.websocket.send_text( + json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) + ) + except Exception: # noqa: BLE001, S110 # client may already be gone + pass + async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): - masked_first: Final = await self._mask_response_create(self.first_message) - self._record_client_frame(masked_first) - await self.backend_ws.send(masked_first) + await self._forward_guarded_frame(self.first_message) while True: message = await self.websocket.receive_text() if not await self._enforce_or_reject_frame(message, self.turn_admission): continue - masked = await self._mask_response_create(message) - self._record_client_frame(masked) - await self.backend_ws.send(masked) + await self._forward_guarded_frame(message) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) @@ -2699,10 +2816,15 @@ class ManagedResponsesWebSocketHandler: first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, turn_admission: Callable[[], Awaitable[None]] | None = None, + frame_guard: ResponsesWsFrameGuard | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, ) -> None: self.websocket = websocket self.model = model + self.frame_guard = frame_guard + self.guardrail_callbacks: tuple[PresidioGuardrailCallback, ...] = tuple(guardrail_callbacks or ()) + self._pii_tokens: Mapping[str, str] = {} self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.litellm_metadata: Mapping[str, object] = litellm_metadata or {} @@ -3011,6 +3133,27 @@ class ManagedResponsesWebSocketHandler: if self.litellm_metadata: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) + async def _guard_frame(self, raw_message: str) -> GuardedResponsesWsFrame | None: + """Same per-frame pre_call guardrail walk as the native socket: a rejected frame errors and is not sent.""" + if self.frame_guard is None: + return GuardedResponsesWsFrame(frame=raw_message) + try: + return await self.frame_guard(raw_message) + except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event + verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) + await self._send_error(str(e), error_type="invalid_request_error") + return None + + @staticmethod + def _metadata_with_turn_guardrails( + existing: object, guardrail_information: tuple[Mapping[str, object], ...] + ) -> Mapping[str, object]: + """Each turn is its own ``litellm.aresponses`` call, so its guardrail records ride on that call's metadata.""" + return { + **(existing if _is_json_object(existing) else {}), + _GUARDRAIL_INFORMATION_KEY: [dict(record) for record in guardrail_information], + } + @staticmethod def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" @@ -3054,12 +3197,19 @@ class ManagedResponsesWebSocketHandler: except Exception: pass try: - await self.websocket.send_text(serialized) + await self.websocket.send_text(self._unmask_for_client(serialized)) except Exception as send_exc: verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc) return terminal_event # Client disconnected return terminal_event + def _unmask_for_client(self, serialized: str) -> str: + """The provider and the turn history keep Presidio's tokens, the client gets the original text back.""" + if not self._pii_tokens or not self.guardrail_callbacks: + return serialized + unmask_pii_text: Final[_UnmasksPiiText] = getattr(self.guardrail_callbacks[0], _UNMASK_PII_TEXT_ATTR) + return _unmask_responses_event(serialized, self._pii_tokens, unmask_pii_text) + def _save_turn_history( self, terminal_event: _MutableJsonObject | None, @@ -3137,7 +3287,13 @@ class ManagedResponsesWebSocketHandler: await _answer_rejected_turn(self.websocket, rejection) return - call_kwargs: Final = self._build_base_call_kwargs(msg_obj) + guarded: Final = await self._guard_frame(raw_message) + if guarded is None: + return + if guarded.pii_tokens: + self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} + + call_kwargs: Final = self._build_base_call_kwargs(_load_json_object(guarded.frame)) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must @@ -3157,6 +3313,10 @@ class ManagedResponsesWebSocketHandler: self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history) self._inject_credentials(call_kwargs, model=model) + if guarded.guardrail_information: + call_kwargs["litellm_metadata"] = self._metadata_with_turn_guardrails( + call_kwargs.get("litellm_metadata"), guarded.guardrail_information + ) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) diff --git a/litellm/router.py b/litellm/router.py index 157fb59d9ba..1411f191863 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6062,6 +6062,9 @@ class Router: including pings and provider error frames, is forwarded live and verbatim instead. """ + from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( aclose_if_supported, aclose_superseded_attempt, @@ -6073,6 +6076,7 @@ class Router: ) source_iterator: Final = response + replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) async def close_source() -> None: with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): @@ -6133,7 +6137,11 @@ class Router: if not has_generated_content and error_event is None else None ) - if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs): + if ( + refusal_stop_details is not None + and not replays_a_finished_agentic_turn + 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, diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 091472c1ea3..b0da3ed264a 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,8 +4,10 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable -from typing import Final, cast +from collections.abc import Callable, Mapping +from typing import Final, TypedDict, cast + +from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.caching.dual_cache import DualCache @@ -16,7 +18,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count -from litellm.llms.bedrock.common_utils import supports_bedrock_prompt_cache_ttl +from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -24,6 +26,11 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p from ..prompt_caching_cache import PromptCachingCache +class _PromptCachingSuccessEventKwargs(TypedDict, total=False): + standard_logging_object: ReadOnly[StandardLoggingPayload | None] + litellm_params: ReadOnly[Mapping[str, object] | None] + + def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int: """ Returns the lowest minimum cacheable prefix across a model group. @@ -49,9 +56,13 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int ) -def _honors_one_hour_cache_ttl(standard_logging_object: StandardLoggingPayload) -> bool: +def _honors_one_hour_cache_ttl( + standard_logging_object: StandardLoggingPayload, litellm_params: Mapping[str, object] +) -> bool: if standard_logging_object["custom_llm_provider"] == "bedrock": - return supports_bedrock_prompt_cache_ttl(standard_logging_object["model"]) + deployment_model: Final = standard_logging_object["model"] + capability_model: Final = bedrock_capability_model(deployment_model, litellm_params.get("base_model")) + return bedrock_cache_point_ttl_supported(capability_model, deployment_model) model_map_value: Final = standard_logging_object["model_map_information"]["model_map_value"] return model_map_value is not None and (model_map_value.get("cache_creation_input_token_cost_above_1hr") or 0) > 0 @@ -121,7 +132,13 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, + kwargs: _PromptCachingSuccessEventKwargs, + response_obj: object, + start_time: object, + end_time: object, + ) -> None: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -163,11 +180,15 @@ class PromptCachingDeploymentCheck(CustomLogger): cache: Final = PromptCachingCache( cache=self.cache, ) + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + raw_litellm_params if isinstance(raw_litellm_params, Mapping) else {} + ) await cache.async_add_model_id( model_id=model_id, messages=prompt_messages, tools=None, # [TODO]: add tools once standard_logging_object supports it - honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object), + honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object, litellm_params), ) return diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index bfd16b29bf6..a82e734b47d 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -151,6 +151,8 @@ class LoggingSurface(Protocol): start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, cache_hit: bool | None = None, + *, + for_async_request: bool = False, **kwargs: object, ) -> None: ... diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index 43b4adc099c..47fec914f9c 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -1,48 +1,9 @@ -import threading -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import Any, Final from openai._models import BaseModel as OpenAIObject -from pydantic import BaseModel, ConfigDict +from pydantic import ConfigDict -from litellm.constants import DEFER_PYDANTIC_BUILD - -_SCHEMA_BUILD_LOCK: Final = threading.RLock() - - -class LiteLLMBaseModel(BaseModel): - model_config = ConfigDict(defer_build=DEFER_PYDANTIC_BUILD) - - if TYPE_CHECKING: - - def __init__(self, /, **data: object) -> None: ... - - @classmethod - def model_rebuild( - cls, - *, - force: bool = False, - raise_errors: bool = True, - _parent_namespace_depth: int = 2, - _types_namespace: Mapping[str, object] | None = None, - ) -> bool | None: - # Resolve names from the model's own module, never a caller frame: a deferred first-use build - # reads f_locals 5 frames up, and on Python < 3.13 that rewrites the dict the caller's locals() returned - with _SCHEMA_BUILD_LOCK: - return super().model_rebuild( - force=force, - raise_errors=raise_errors, - _parent_namespace_depth=0, - _types_namespace=_types_namespace, - ) - - def model_post_init(self, context: object, /) -> None: - # Instances built by a parent's validator or by model_construct skip this class's own - # first-use build, leaving a MockValSer that fails when the instance is serialized as Any - # TODO: drop once pydantic ships a fix for https://github.com/pydantic/pydantic/issues/13647 - model = type(self) - if not model.__pydantic_complete__: - model.model_rebuild(raise_errors=False) +from litellm.types.llms.base_model import LiteLLMBaseModel class LiteLLMPydanticObjectBase(LiteLLMBaseModel): diff --git a/litellm/types/llms/base_model.py b/litellm/types/llms/base_model.py new file mode 100644 index 00000000000..44088b7a4f1 --- /dev/null +++ b/litellm/types/llms/base_model.py @@ -0,0 +1,44 @@ +import threading +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict + +from litellm.constants import DEFER_PYDANTIC_BUILD + +_SCHEMA_BUILD_LOCK: Final = threading.RLock() + + +class LiteLLMBaseModel(BaseModel): + model_config = ConfigDict(defer_build=DEFER_PYDANTIC_BUILD) + + if TYPE_CHECKING: + + def __init__(self, /, **data: object) -> None: ... + + @classmethod + def model_rebuild( + cls, + *, + force: bool = False, + raise_errors: bool = True, + _parent_namespace_depth: int = 2, + _types_namespace: Mapping[str, object] | None = None, + ) -> bool | None: + # Resolve names from the model's own module, never a caller frame: a deferred first-use build + # reads f_locals 5 frames up, and on Python < 3.13 that rewrites the dict the caller's locals() returned + with _SCHEMA_BUILD_LOCK: + return super().model_rebuild( + force=force, + raise_errors=raise_errors, + _parent_namespace_depth=0, + _types_namespace=_types_namespace, + ) + + def model_post_init(self, context: object, /) -> None: + # Instances built by a parent's validator or by model_construct skip this class's own + # first-use build, leaving a MockValSer that fails when the instance is serialized as Any + # TODO: drop once pydantic ships a fix for https://github.com/pydantic/pydantic/issues/13647 + model = type(self) + if not model.__pydantic_complete__: + model.model_rebuild(raise_errors=False) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 054a0a1c034..8e56ebdc39b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -204,6 +204,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_pdf_input: bool | None supports_native_streaming: bool | None supports_native_structured_output: bool | None + supports_bedrock_native_structured_output: ReadOnly[bool | None] supports_parallel_function_calling: bool | None supports_web_search: bool | None supports_reasoning: bool | None @@ -4047,7 +4048,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( - {"supports_regex_lookaround", "supports_assistant_prefill"} + {"supports_regex_lookaround", "supports_assistant_prefill", "supports_bedrock_native_structured_output"} ) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( @@ -4269,6 +4270,7 @@ class LlmProviders(str, Enum): DEEPINFRA = "deepinfra" PERPLEXITY = "perplexity" TYPESAFE = "typesafe" + NACE = "nace" STRANDS_DECIDER = "strands_decider" LAYA = "laya" BESPOKE = "bespoke" diff --git a/litellm/utils.py b/litellm/utils.py index 3b43c570466..b37827a3b13 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -44,7 +44,7 @@ from httpx import Proxy from httpx._utils import get_environment_proxies from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # OpenAI parser module is private from openai.types.chat.completion_create_params import ResponseFormat -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, with_config +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config import litellm import litellm.litellm_core_utils @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -3303,46 +3304,176 @@ _BEDROCK_REGION_PREFIXES: Final = ( "ap-northeast-1.", ) -_CACHE_PRICING_FIELDS: Final = ( - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost", - "cache_read_input_token_cost_above_200k_tokens", +_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token") + + +def _builtin_cost_map() -> dict[str, Mapping[str, object]]: + return cast(dict[str, Mapping[str, object]], litellm.model_cost) + + +def _replace_cost_map_entry(cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object]) -> None: + current: Final = cost_map.get(key) + merged: Final = {**(current or {}), **entry} + if not isinstance(current, AliasedCostEntry): + cost_map.update({key: merged}) + return + sharing_names: Final = tuple(name for name in current.names if cost_map.get(name) is current) + cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, current.names))) + + +def _without_synthesized_costs( + model_info: Mapping[str, object], key: str, value: Mapping[str, object] +) -> Mapping[str, object]: + """``_get_model_info_helper`` synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` = 0 when the raw + entry has neither (#30198). Persisting those zeros would turn a name-priced entry into a free one and make + ``_is_cost_explicitly_configured`` disable budget enforcement on the next re-registration.""" + builtin_cost_map: Final = _builtin_cost_map() + raw_fields: Final = frozenset(builtin_cost_map.get(str(model_info["key"]), builtin_cost_map.get(key, {}))) + return MappingProxyType( + { + name: field_value + for name, field_value in model_info.items() + if name not in _TOKEN_COST_FIELDS or name in raw_fields or name in value + } + ) + + +_DIRECTION_RATE_FIELDS: Final = MappingProxyType( + { + "input_cost_per_token": ("input_cost_per_token",), + "output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"), + } ) +_RATE_TABLE: Final = TypeAdapter(Mapping[str, object]) +_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) -def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, object] | None: +def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool: + is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields)) + return name in overridden_fields or (is_variant and "_cache_" not in name) + + +def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]: + return MappingProxyType( + {name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)} + ) + + +def _rate_table_with_flat_prices( + name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object] +) -> object: + try: + if name == "tiered_pricing": + return [ + {**_without_overridden_rates(tier, overridden_fields), **flat_prices} + for tier in _RATE_TIERS.validate_python(table) + ] + if name == "off_peak_pricing": + return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields)) + except ValidationError: + return table + return table + + +def _without_overridden_price_variants( + entry: Mapping[str, object], value: Mapping[str, object] +) -> Mapping[str, object]: + """A caller's flat ``input_cost_per_token`` / ``output_cost_per_token`` is that direction's price above every + token threshold, for every service tier and batch, in every ``tiered_pricing`` tier and off-peak window, so each + inherited rate that would stand in for it there, the reasoning-token rate included for output, is dropped, and + each tier carries the flat price instead.""" + flat_prices: Final = MappingProxyType( + {cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None} + ) + overridden_fields: Final = tuple( + itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices) + ) + if not overridden_fields: + return entry + return MappingProxyType( + { + name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices) + for name, field_value in entry.items() + if not _is_overridden_rate(name, overridden_fields) + } + ) + + +@dataclass(frozen=True, slots=True) +class _BuiltinCostMatch: + key: str + entry: Mapping[str, object] + names_the_row: bool + + +@dataclass(frozen=True, slots=True) +class _LookupMiss: + tripped: bool + + +@dataclass(frozen=True, slots=True) +class _FallbackRegistration: + entry: Mapping[str, object] + matched_builtin: bool + + +def _fallback_registration(key: str, provider: str, value: Mapping[str, object]) -> _FallbackRegistration: + """The key keeps its own entry, seeded from the capability rules an unregistered lookup resolves it to, with the + matched built-in entry on top. The matched row's provider stays behind, since the key's own lookup can reject it + (an ``azure_ai/`` key never resolves to an ``openai`` entry), and the key keeps the provider its caller gave.""" + match: Final = _resolve_builtin_model_cost_entry(key=key, provider=provider) + generalized: Final[Mapping[str, object]] = match_capability_generalizations(key) or {} + if match is None: + return _FallbackRegistration(entry=generalized, matched_builtin=False) + inherited: Final = match.entry if match.names_the_row else _without_overridden_price_variants(match.entry, value) + own_fields: Final = {name: field_value for name, field_value in inherited.items() if name != "litellm_provider"} + return _FallbackRegistration(entry={**generalized, **own_fields}, matched_builtin=True) + + +def _resolve_builtin_model_cost_entry(key: str, provider: str) -> _BuiltinCostMatch | None: """Best-effort lookup of a built-in ``model_cost`` entry for a custom key whose shape ``get_model_info`` cannot resolve (repeated provider prefixes like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region - aliases). + aliases), or whose lookup tripped. - Returns a copy of the matching entry so the caller can inherit its defaults - (most importantly cache pricing) without mutating the shared built-in. - Returns ``None`` when no safe match exists. + The key itself is tried first, then the key behind its provider prefixes + and bedrock route, so a catalog row whose lookup tripped inherits from + itself rather than from a region-stripped sibling priced at the base rate. + ``names_the_row`` is set when the match is the row the key names, which + ``get_model_info`` resolves it to, so the caller keeps that row's rates as + an untripped registration does. Returns ``None`` when no candidate exists + under a provider ``get_model_info`` would accept for ``provider``, or, when + the value names none, for the provider prefix the key itself carries, + which is what ``get_model_info`` derives when it reads the key back. """ - candidates: Final[list[str]] = [] + from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix + segments: Final = key.split("/") - idx = 0 - while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet: - idx += 1 - candidates.append("/".join(segments[idx:])) - - base: Final = candidates[-1] if candidates else key - for region_prefix in _BEDROCK_REGION_PREFIXES: - if base.startswith(region_prefix): - candidates.append(base[len(region_prefix) :]) - - if provider: - stripped: Final = _strip_model_name(model=base, custom_llm_provider=provider) - if stripped != base: - candidates.append(stripped) + match_provider: Final = provider or (segments[0] if len(segments) > 1 and segments[0] in LlmProvidersSet else "") + provider_prefix_count: Final = next( + (idx for idx, segment in enumerate(segments[:-1]) if segment not in LlmProvidersSet), len(segments) - 1 + ) + prefix_stripped: Final = tuple("/".join(segments[idx:]) for idx in range(1, provider_prefix_count + 1)) + prefix_free: Final = prefix_stripped[-1] if prefix_stripped else key + base: Final = ( + strip_bedrock_routing_prefix(prefix_free) if match_provider in ("bedrock", "bedrock_converse") else prefix_free + ) + region_stripped: Final = tuple( + base[len(region_prefix) :] for region_prefix in _BEDROCK_REGION_PREFIXES if base.startswith(region_prefix) + ) + stripped: Final = _strip_model_name(model=base, custom_llm_provider=match_provider) if match_provider else base + row_names: Final = (key, *prefix_stripped, base, stripped) + candidates: Final = (key, *prefix_stripped, base, *region_stripped, stripped) + builtin_cost_map: Final = _builtin_cost_map() for candidate in candidates: - entry = litellm.model_cost.get(candidate) - if entry is not None and entry.get("litellm_provider") is not None: - return dict(entry) + entry = builtin_cost_map.get(candidate) + if ( + entry is not None + and entry.get("litellm_provider") is not None + and _check_provider_match(dict(entry), match_provider) + ): + return _BuiltinCostMatch(key=candidate, entry=dict(entry), names_the_row=candidate in row_names) return None @@ -3360,19 +3491,115 @@ def is_generalized_model_info(model_info: ModelInfo) -> bool: return key not in litellm.model_cost and match_capability_generalizations(key) is not None -def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | None: - """Resolve ``model`` to its built-in cost-map entry for registration merging. +_REGISTRATION_ATTEMPTS: Final = 3 - Returns ``None`` when the lookup raises or when it resolved via a - fallback-generalization capability rule. A rule-derived entry carries no - pricing, so treating it as a hit would skip the built-in cache-pricing - inheritance for prefix-mangled keys. - """ + +def _lookup_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss: try: info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except ModelNotMappedError: + return _LookupMiss(tripped=False) except Exception: - return None - return None if is_generalized_model_info(info) else info + return _LookupMiss(tripped=True) + return _LookupMiss(tripped=False) if is_generalized_model_info(info) else info + + +def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss: + """Resolve ``model`` to its built-in cost-map entry for registration merging. + + Misses when the lookup cannot map the model or resolved it via a + fallback-generalization capability rule, since a rule-derived entry carries no + pricing. A lookup that raised anything else, such as a concurrent registration + changing the map mid-walk, is retried, so a transient failure registers exactly + where the untripped lookup would. + """ + lookups: Final = (_lookup_for_registration(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS)) + return next( + (lookup for lookup in lookups if not (isinstance(lookup, _LookupMiss) and lookup.tripped)), + _LookupMiss(tripped=True), + ) + + +@dataclass(frozen=True, slots=True) +class _RegistrationWrite: + key: str + entry: Mapping[str, object] + + +def _registration_base( + key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None +) -> _RegistrationWrite: + provider: Final = str(value.get("litellm_provider") or "") + if provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO or any( + key.startswith(f"{p}/") for p in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO + ): + return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {}))) + builtin_model_info: Final = _get_builtin_model_info_for_registration( + model=key, custom_llm_provider=custom_llm_provider + ) + if not isinstance(builtin_model_info, _LookupMiss): + return _RegistrationWrite( + key=builtin_model_info["key"], + entry=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)), + ) + fallback: Final = _fallback_registration(key=key, provider=provider, value=value) + if not fallback.matched_builtin and ( + value.get("cache_creation_input_token_cost") is None + and value.get("cache_read_input_token_cost") is None + and value.get("tiered_pricing") is None + and (value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None) + ): + verbose_logger.warning( + "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info", + warning_display_name or key, + ) + return _RegistrationWrite(key=key, entry=dict(fallback.entry)) + + +def _registration_write( + key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None +) -> _RegistrationWrite: + base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name) + # ``get_model_info`` returns ``litellm_provider: None`` when the + # provider is unknown (e.g. custom deployments registered via + # ``Router.add_deployment``). Persisting that None into + # ``litellm.model_cost`` causes ``_check_provider_match`` to drop + # custom pricing on subsequent cost lookups. + existing_model: Final = { + name: field_value + for name, field_value in base.entry.items() + if name != "litellm_provider" or field_value is not None + } + return _RegistrationWrite(key=base.key, entry=update_dictionary(existing_model, value)) + + +def _registration_attempt( + key: str, + value: Mapping[str, object], + custom_llm_provider: str | None, + warning_display_name: str | None, + final_attempt: bool, +) -> _RegistrationWrite | None: + source: Final = _builtin_cost_map() + write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name) + _replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry) + return write if final_attempt or litellm.model_cost is source else None + + +def _register_on_current_map( + key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None +) -> _RegistrationWrite: + """A price data reload that swaps the map while the entry is being built or written would otherwise get the old + map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was + built from and the entry is rebuilt when a reload replaced that map before the write finished. The last attempt + writes into the current map, so the entry is never lost.""" + attempts: Final = ( + _registration_attempt( + key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1 + ) + for attempt in range(_REGISTRATION_ATTEMPTS) + ) + return next(write for write in attempts if write is not None) _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload @@ -3492,79 +3719,15 @@ def register_model( for _registered_key, _registered_value in _registrations.items(): _runtime_registered_model_cost[_registered_key] = dict(_registered_value) - _skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO - for key, value in loaded_model_cost.items(): - ## get model info ## - provider = value.get("litellm_provider", "") - _key_str = str(key) - if provider in _skip_get_model_info_providers or any( - _key_str.startswith(f"{p}/") for p in _skip_get_model_info_providers - ): - existing_model = litellm.model_cost.get(key, {}) - model_cost_key = key - else: - builtin_model_info = _get_builtin_model_info_for_registration( - model=_key_str, custom_llm_provider=custom_llm_provider - ) - if builtin_model_info is not None: - existing_model = cast(dict, builtin_model_info) - model_cost_key = existing_model["key"] - else: - # An exact entry ends the lookup ladder before the capability rules are - # consulted, so seed from them: otherwise registering an unmapped model - # shadows the very defaults it would have resolved to unregistered. - existing_model = dict(match_capability_generalizations(_key_str) or {}) - model_cost_key = key - builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider) - if builtin_entry is not None: - for field in _CACHE_PRICING_FIELDS: - if value.get(field) is None and builtin_entry.get(field) is not None: - existing_model[field] = builtin_entry[field] - elif ( - value.get("cache_creation_input_token_cost") is None - and value.get("cache_read_input_token_cost") is None - and value.get("tiered_pricing") is None - and ( - value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None - ) - ): - verbose_logger.warning( - "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info", - warning_display_name or key, - ) - # ``get_model_info`` returns ``litellm_provider: None`` when the - # provider is unknown (e.g. custom deployments registered via - # ``Router.add_deployment``). Persisting that None into - # ``litellm.model_cost`` causes ``_check_provider_match`` to drop - # custom pricing on subsequent cost lookups. - if existing_model.get("litellm_provider") is None: - existing_model.pop("litellm_provider", None) - # Same pattern for cost fields (#30198): ``_get_model_info_helper`` - # synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` - # = 0 when they are absent from the raw entry. Writing those zeros - # back flips a sparse entry from "no cost keys" (priced via name) - # to "cost keys = 0" (free), which makes - # ``_is_cost_explicitly_configured`` return True and silently - # disables budget enforcement on the next re-registration. - _raw_entry = litellm.model_cost.get(model_cost_key) - if _raw_entry is None: - _raw_entry = litellm.model_cost.get(key) - if _raw_entry is None: - _raw_entry = {} - for _cost_field in ("input_cost_per_token", "output_cost_per_token"): - if _cost_field not in _raw_entry and _cost_field not in value: - existing_model.pop(_cost_field, None) - ## override / add new keys to the existing model cost dictionary - updated_dictionary = update_dictionary(existing_model, value) - litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) + write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name) # Invalidate case-insensitive lookup map since model_cost was modified _invalidate_model_cost_lowercase_map() - verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", model_cost_key, model_cost_key) + verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", write.key, write.key) if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter": - litellm.openrouter_models.add(_key_str.split("/", 1)[-1]) + litellm.openrouter_models.add(str(key).split("/", 1)[-1]) if add_to_provider_model_sets: litellm.add_known_models(loaded_model_cost) return model_cost @@ -4575,6 +4738,7 @@ def get_optional_params( store: bool | None = None, prompt_cache_key: str | None = None, base_model: str | None = None, + model_info: Mapping[str, object] | None = None, **kwargs: object, ): drop_params = normalize_drop_params(drop_params) # rebind-ok: config and DB deployments pass "true" as a string @@ -4585,6 +4749,7 @@ def get_optional_params( # non_default_params / _check_valid_arg — it's a routing hint, not an # OpenAI param. passed_params.pop("base_model", None) + passed_params.pop("model_info", None) provider_config: BaseConfig | None = None if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: provider_config = ProviderConfigManager.get_provider_chat_config( @@ -4824,6 +4989,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, drop_params=bool(drop_params), + model_info=model_info, ) elif bedrock_route == "openai": optional_params = litellm.AmazonBedrockOpenAIConfig().map_openai_params( @@ -4841,7 +5007,7 @@ def get_optional_params( drop_params=bool(drop_params), ) else: - optional_params = litellm.AmazonAnthropicClaudeConfig().map_openai_params( + optional_params = litellm.AmazonAnthropicClaudeConfig(model_info=model_info).map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, @@ -5729,7 +5895,7 @@ def _rebuild_model_cost_lowercase_map() -> dict[str, str]: The rebuilt map (guaranteed to be not None). """ global _model_cost_lowercase_map - _model_cost_lowercase_map = {k.lower(): k for k in litellm.model_cost} + _model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)} return _model_cost_lowercase_map @@ -5744,9 +5910,7 @@ def _handle_stale_map_entry_rebuild( Returns: The matched key if found after rebuild, None otherwise. """ - global _model_cost_lowercase_map - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() - matched_key: Final = _model_cost_lowercase_map.get(potential_key_lower) + matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower) if matched_key is not None and matched_key in litellm.model_cost: return matched_key return None @@ -5764,7 +5928,7 @@ def _handle_new_key_with_scan( The matched key if found, None otherwise. """ global _model_cost_lowercase_map - for key in litellm.model_cost: + for key in tuple(_builtin_cost_map()): if key.lower() == potential_key_lower: _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() return key @@ -5786,18 +5950,19 @@ def get_model_cost_key(potential_key: str) -> str | None: called and confirmed not to cause performance issues, add it to the allowed_helpers list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py """ - global _model_cost_lowercase_map - # Exact match (O(1)) if potential_key in litellm.model_cost: return potential_key - # Case-insensitive lookup via map (O(1)) - if _model_cost_lowercase_map is None: - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() + # Case-insensitive lookup via map (O(1)); bound locally because another thread's + # registration can invalidate the module-level map at any point + cached_lowercase_map: Final = _model_cost_lowercase_map + lowercase_map: Final = ( + cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map() + ) potential_key_lower: Final = potential_key.lower() - matched_key = _model_cost_lowercase_map.get(potential_key_lower) + matched_key = lowercase_map.get(potential_key_lower) # Verify key exists (O(1) - handles model_cost.pop() case) if matched_key is not None and matched_key in litellm.model_cost: @@ -5937,6 +6102,15 @@ def _get_model_info_from_generalization( return None +def unmapped_model_generalization(model: str, custom_llm_provider: str | None) -> Mapping[str, object] | None: + generalization: Final = _get_model_info_from_generalization( + model=model, + potential_model_names=get_potential_model_names(model, custom_llm_provider), + custom_llm_provider=custom_llm_provider, + ) + return generalization[1] if generalization is not None else None + + def _strip_mantle_region_prefix(model: str) -> str: from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix @@ -6465,6 +6639,9 @@ def get_model_info_helper( supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None), supports_native_streaming=_model_info.get("supports_native_streaming", None), supports_native_structured_output=_model_info.get("supports_native_structured_output", None), + supports_bedrock_native_structured_output=_model_info.get( + "supports_bedrock_native_structured_output", None + ), supports_web_search=_model_info.get("supports_web_search", None), supports_url_context=_model_info.get("supports_url_context", None), supports_reasoning=_model_info.get("supports_reasoning", None), @@ -9103,6 +9280,8 @@ class ProviderConfigManager: return litellm.PerplexityDecisionsConfig() if provider == LlmProviders.TYPESAFE: return litellm.TypeSafeDecisionsConfig() + if provider == LlmProviders.NACE: + return litellm.NaceDecisionsConfig() if provider == LlmProviders.OPENROUTER: return litellm.OpenRouterDecisionsConfig() if provider == LlmProviders.FIREWORKS_AI: diff --git a/litellm_cli/__init__.py b/litellm_cli/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm_cli/main.py b/litellm_cli/main.py new file mode 100644 index 00000000000..7b23d0a5bb5 --- /dev/null +++ b/litellm_cli/main.py @@ -0,0 +1,183 @@ +"""Console entry points for `lite` and `litellm-proxy`. + +The CLI is a thin HTTP client, but it lives inside the `litellm` package, so a plain +import would run the whole SDK `__init__` before Click parses argv. These entry points +import the CLI against an uninitialized `litellm` package object instead. CLI code under +`litellm.proxy.client`, and whatever it imports, loads without the SDK whenever it is +first imported. The first attribute read on the SDK, or the first import of any other +new code under `litellm`, initializes the SDK in place before that code runs, so code +that needs the SDK sees it exactly as a plain `import litellm` would leave it. An init that +raises leaves the SDK uninitialized, so the next use runs it again and raises the same error. +""" + +import builtins +import contextvars +import importlib +import importlib.util +import sys +import threading +import types +from collections.abc import Callable, Mapping, Sequence +from typing import Final + +SDK_PACKAGE: Final = "litellm" +CLI_PACKAGE: Final = "litellm.proxy.client" +_INITIALIZATION_LOCK: Final = threading.RLock() +_BUILTIN_IMPORT: Final = builtins.__import__ +_IMPORT_MODULE: Final = importlib.import_module +_LOADING_CLI_CODE: Final = contextvars.ContextVar("_LOADING_CLI_CODE", default=False) + + +def _is_dunder(name: str) -> bool: + return name.startswith("__") and name.endswith("__") + + +class _SDKUntilFirstUse(types.ModuleType): + def __getattribute__(self, name: str) -> object: + if not _is_dunder(name): + _initialize(self) + return types.ModuleType.__getattribute__(self, name) + + +class _SDKInitializing(types.ModuleType): + def __getattr__(self, name: str) -> object: + with _INITIALIZATION_LOCK: + if type(self) is _SDKInitializing: + raise AttributeError(f"partially initialized module {SDK_PACKAGE!r} has no attribute {name!r}") + return getattr(self, name) + + +def _become(sdk: types.ModuleType, state: type[types.ModuleType]) -> None: + sdk.__class__ = state # rebind-ok: every holder of the litellm module must see its state change in place + + +def _restore_interpreter_imports() -> None: + if builtins.__import__ is _import_initializing_sdk: + builtins.__import__ = _BUILTIN_IMPORT + if importlib.import_module is _import_module_initializing_sdk: + importlib.import_module = _IMPORT_MODULE + + +def _initialize(sdk: types.ModuleType) -> None: + with _INITIALIZATION_LOCK: + if type(sdk) is not _SDKUntilFirstUse: + return + spec: Final = sdk.__spec__ + if spec is None or spec.loader is None: + raise ModuleNotFoundError(f"No module named {SDK_PACKAGE!r}", name=SDK_PACKAGE) + _become(sdk, _SDKInitializing) + try: + spec.loader.exec_module(sdk) + except BaseException: + _become(sdk, _SDKUntilFirstUse) + raise + _become(sdk, types.ModuleType) + _restore_interpreter_imports() + + +def _in_package(target: str, package: str) -> bool: + return target == package or target.startswith(f"{package}.") + + +def _needs_more_of(loaded: types.ModuleType, name: str) -> bool: + if name == "*": + return type(loaded) is _SDKUntilFirstUse + return name not in vars(loaded) + + +def _loads_new_sdk_code(target: str, fromlist: Sequence[str]) -> bool: + if not _in_package(target, SDK_PACKAGE): + return False + loaded: Final = sys.modules.get(target) + if loaded is None: + return True + return any(_needs_more_of(loaded, name) for name in fromlist) + + +def _load_cli_code(load: Callable[[], types.ModuleType]) -> types.ModuleType: + token: Final = _LOADING_CLI_CODE.set(True) + try: + return load() + finally: + _LOADING_CLI_CODE.reset(token) + + +def _imported(target: str | None, fromlist: Sequence[str], load: Callable[[], types.ModuleType]) -> types.ModuleType: + if _LOADING_CLI_CODE.get() or target is None or not _loads_new_sdk_code(target, fromlist): + return load() + if _in_package(target, CLI_PACKAGE): + return _load_cli_code(load) + sdk: Final = sys.modules.get(SDK_PACKAGE) + if sdk is not None: + _initialize(sdk) + return load() + + +def _import_target(name: str, importer_globals: Mapping[str, object] | None, level: int) -> str | None: + if level == 0: + return name + package: Final = (importer_globals or {}).get("__package__") + if not isinstance(package, str) or not package: + return None + return importlib.util.resolve_name("." * level + name, package) + + +def _import_initializing_sdk( + name: str, + globals: Mapping[str, object] | None = None, + locals: Mapping[str, object] | None = None, + fromlist: Sequence[str] | None = (), + level: int = 0, +) -> types.ModuleType: + return _imported( + _import_target(name, globals, level), + fromlist or (), + lambda: _BUILTIN_IMPORT(name, globals, locals, fromlist or (), level), + ) + + +def _import_module_target(name: str, package: str | None) -> str | None: + if not name.startswith("."): + return name + if not package: + return None + return importlib.util.resolve_name(name, package) + + +def _import_module_initializing_sdk(name: str, package: str | None = None) -> types.ModuleType: + return _imported(_import_module_target(name, package), (), lambda: _IMPORT_MODULE(name, package)) + + +def _install_sdk_until_first_use() -> None: + sdk_spec: Final = importlib.util.find_spec(SDK_PACKAGE) + if sdk_spec is None: + raise ModuleNotFoundError(f"No module named {SDK_PACKAGE!r}", name=SDK_PACKAGE) + sdk: Final = importlib.util.module_from_spec(sdk_spec) + sdk.__class__ = _SDKUntilFirstUse + sys.modules[SDK_PACKAGE] = sdk + + from litellm._process_defaults import apply_process_defaults + + apply_process_defaults() + builtins.__import__ = _import_initializing_sdk + importlib.import_module = _import_module_initializing_sdk + + +def _entry_points() -> tuple[Callable[[], None], Callable[[], None]]: + if SDK_PACKAGE not in sys.modules: + _install_sdk_until_first_use() + + from litellm.proxy.client.cli.main import cli as lite + from litellm.proxy.client.cli.main import litellm_proxy_cli as deprecated_litellm_proxy + + return lite, deprecated_litellm_proxy + + +def cli() -> None: + lite, _ = _entry_points() + lite() + + +def litellm_proxy_cli() -> None: + _, deprecated_litellm_proxy = _entry_points() + deprecated_litellm_proxy() diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 255b94b0ea8..b39a9bfa9bf 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -24,6 +24,32 @@ FROM $UV_IMAGE AS uvbin # under nodeenv. Crucially we do NOT run `prisma generate` at RUNTIME; the # old migration job did, on every pod start, which is the wasteful behaviour # the componentization is fixing. +# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes +FROM $LITELLM_BUILD_IMAGE AS rust-native + +WORKDIR /app +USER root + +COPY --from=uvbin /uv /uvx /usr/local/bin/ + +RUN for i in 1 2 3; do \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done + +ENV UV_PYTHON_DOWNLOADS=0 + +COPY pyproject.toml README.md rust-toolchain.toml ./ +COPY .cargo ./.cargo +COPY litellm-rust ./litellm-rust + +RUN --mount=type=cache,target=/root/.cache/uv \ + mkdir -p litellm/rust_bridge /wheels && \ + touch litellm/__init__.py litellm/rust_bridge/__init__.py && \ + uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ + python -c "import maturin; maturin.build_wheel('/wheels')" + FROM $LITELLM_BUILD_IMAGE AS builder WORKDIR /app @@ -58,15 +84,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --python python3.13 -# Stage 2 — copy source and install the project + workspace members. +# Stage 2 — copy source and install the workspace members. COPY . . RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-default-groups --no-editable \ + uv sync --frozen --no-install-project --no-default-groups --no-editable \ --extra proxy \ --extra extra_proxy \ --python python3.13 +RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \ + uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl + RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ COPY migrations/run.py /app/run.py diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c3efaf168ab..6ebb54c0035 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1546,7 +1546,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -1623,7 +1623,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -1700,7 +1700,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -1777,7 +1777,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -12877,6 +12877,7 @@ "supports_prompt_caching": true }, "azure_ai/ministral-3b": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-08, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12940,6 +12941,7 @@ "supports_vision": true }, "azure_ai/mistral-medium-2505": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12977,6 +12979,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -29699,6 +29702,7 @@ }, "vertex_ai/gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -29806,6 +29810,7 @@ "supports_native_streaming": true }, "vertex_ai/gemini-3.8-live": { + "deprecation_date": "2027-09-25", "input_cost_per_audio_token": 3e-06, "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, @@ -31960,6 +31965,7 @@ }, "gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -32409,11 +32415,16 @@ ] }, "github_copilot/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32427,11 +32438,16 @@ "thinking_always_on": true }, "github_copilot/claude-fable-5.1": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32490,11 +32506,16 @@ "thinking_always_on": false }, "github_copilot/claude-opus-4.8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32524,13 +32545,17 @@ "supports_vision": true }, "github_copilot/claude-opus-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, "deprecation_date": "2026-11-16", + "input_cost_per_token": 5e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", + "output_cost_per_token": 2.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32543,11 +32568,16 @@ "supports_vision": true }, "github_copilot/claude-opus-5.5": { + "cache_creation_input_token_cost": 5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 2e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32561,11 +32591,16 @@ "thinking_always_on": true }, "github_copilot/claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32578,11 +32613,16 @@ "supports_vision": true }, "github_copilot/claude-sonnet-5.5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32596,11 +32636,15 @@ "thinking_always_on": true }, "github_copilot/gemini-3.7-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions" ], @@ -32610,11 +32654,15 @@ "supports_vision": true }, "github_copilot/gemini-3.8-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions" ], @@ -32780,11 +32828,18 @@ "supports_vision": true }, "github_copilot/gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -32796,11 +32851,15 @@ "supports_vision": true }, "github_copilot/gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 4.5e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32811,11 +32870,18 @@ "supports_vision": true }, "github_copilot/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32826,11 +32892,20 @@ "supports_vision": true }, "github_copilot/gpt-5.6-luna": { + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_200k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_200k_tokens": 4e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_200k_tokens": 4e-07, "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32841,11 +32916,20 @@ "supports_vision": true }, "github_copilot/gpt-5.6-sol": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32856,13 +32940,21 @@ "supports_vision": true }, "github_copilot/gpt-5.6-terra": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, "deprecation_date": "2026-11-16", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32873,11 +32965,20 @@ "supports_vision": true }, "github_copilot/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32888,11 +32989,20 @@ "supports_vision": true }, "github_copilot/gpt-6-luna": { + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "input_cost_per_token": 1e-07, + "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 5e-07, + "output_cost_per_token_above_272k_tokens": 7.5e-07, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32903,11 +33013,20 @@ "supports_vision": true }, "github_copilot/gpt-6-sol": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32918,11 +33037,20 @@ "supports_vision": true }, "github_copilot/gpt-6.1-sol": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_above_272k_tokens": 2e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -32933,11 +33061,15 @@ "supports_vision": true }, "github_copilot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, "litellm_provider": "github_copilot", "max_input_tokens": 917504, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/chat/completions" ], @@ -32947,11 +33079,15 @@ "supports_vision": true }, "github_copilot/mai-code-1.1-flash": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, "litellm_provider": "github_copilot", "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "output_cost_per_token": 1.2e-06, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", "supported_endpoints": [ "/v1/responses" ], @@ -45295,6 +45431,38 @@ "output_cost_per_token": 0.0, "source": "https://docs.perplexity.ai/docs/getting-started/pricing" }, + "perplexity/pplx-decider-v1.2": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-latest": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-v1.2": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-latest": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, "perplexity/sonar": { "input_cost_per_token": 1e-06, "litellm_provider": "perplexity", @@ -65431,6 +65599,40 @@ "supports_vision": false, "source": "https://deepinfra.com/pricing" }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash" + }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash-1017" + }, "deepinfra/openai/gpt-oss-120b-Ultra": { "max_tokens": 131072, "max_input_tokens": 131072, @@ -71184,6 +71386,38 @@ "notes": "Self-hosted decision model; infrastructure costs are paid separately" } }, + "nace/drex-v1.0": { + "input_cost_per_token": 4e-08, + "litellm_provider": "nace", + "max_input_tokens": 32768, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v1.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v2.0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-latest": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", @@ -72466,7 +72700,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72499,7 +72733,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72856,6 +73090,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72887,6 +73122,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72933,6 +73169,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73003,6 +73240,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -73039,6 +73277,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -73082,6 +73321,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73123,6 +73363,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73164,6 +73405,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -73286,6 +73528,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -73329,6 +73572,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74301,6 +74545,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74337,6 +74582,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74388,6 +74634,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74468,6 +74715,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -74509,6 +74757,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -74557,6 +74806,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74603,6 +74853,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74649,6 +74900,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -74699,6 +74951,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -74747,6 +75000,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -75141,6 +75395,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-4-argon": { + "deprecation_date": "2027-10-13", "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 4404aa8e35c..33a784e2dae 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1027,6 +1027,9 @@ "supports_bedrock_converse_tool_result_images": { "type": "boolean" }, + "supports_bedrock_native_structured_output": { + "type": "boolean" + }, "supports_bedrock_runtime_chat_completions_inline_reasoning": { "type": "boolean" }, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 0325fe1f796..ad38b064903 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2625,6 +2625,13 @@ "image_variations": true } }, + "nace": { + "display_name": "Nace (`nace`)", + "url": "https://console.nace.ai/docs/drex", + "endpoints": { + "systemone": true + } + }, "typesafe": { "display_name": "TypeSafe (`typesafe`)", "url": "https://docs.typesafe.ai/models", diff --git a/pyproject.toml b/pyproject.toml index 854811799a0..9262630f63e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -190,8 +190,8 @@ proxy-runtime = [ [project.scripts] litellm = "litellm:run_server" -lite = "litellm.proxy.client.cli:cli" -litellm-proxy = "litellm.proxy.client.cli:litellm_proxy_cli" +lite = "litellm_cli.main:cli" +litellm-proxy = "litellm_cli.main:litellm_proxy_cli" [dependency-groups] dev = [ @@ -312,6 +312,7 @@ build-backend = "maturin" manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml" module-name = "litellm.rust_bridge._native" python-source = "." +python-packages = ["litellm_cli"] bindings = "pyo3" features = ["extension-module"] profile = "release" diff --git a/pyrightconfig.json b/pyrightconfig.json index 2264c29be37..3a0812413c3 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -1,5 +1,5 @@ { - "include": ["litellm"], + "include": ["litellm", "litellm_cli"], "ignore": [], "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "tests/e2e_harness/claude_code", "tests/e2e/ui", "litellm/types/utils.py", "litellm/proxy/_types.py"], "executionEnvironments": [ diff --git a/terraform/litellm/aws/.terraform.lock.hcl b/terraform/litellm/aws/.terraform.lock.hcl index 30b2a194c4c..79097043254 100644 --- a/terraform/litellm/aws/.terraform.lock.hcl +++ b/terraform/litellm/aws/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/aws" { constraints = "~> 5.60" hashes = [ "h1:Ijt7pOlB7Tr7maGQIqtsLFbl7pSMIj06TVdkoSBcYOw=", + "h1:edXOJWE4ORX8Fm+dpVpICzMZJat4AX0VRCAy/xkcOc0=", "zh:054b8dd49f0549c9a7cc27d159e45327b7b65cf404da5e5a20da154b90b8a644", "zh:0b97bf8d5e03d15d83cc40b0530a1f84b459354939ba6f135a0086c20ebbe6b2", "zh:1589a2266af699cbd5d80737a0fe02e54ec9cf2ca54e7e00ac51c7359056f274", @@ -28,6 +29,7 @@ provider "registry.terraform.io/hashicorp/random" { version = "3.8.1" constraints = "~> 3.6" hashes = [ + "h1:Eexl06+6J+s75uD46+WnZtpJZYRVUMB0AiuPBifK6Jc=", "h1:u8AKlWVDTH5r9YLSeswoVEjiY72Rt4/ch7U+61ZDkiQ=", "zh:08dd03b918c7b55713026037c5400c48af5b9f468f483463321bd18e17b907b4", "zh:0eee654a5542dc1d41920bbf2419032d6f0d5625b03bd81339e5b33394a3e0ae", diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 6fcbdc2f500..bedb6f4e98b 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -419,8 +419,9 @@ To run *many* tenants from a single config, call the module with ```hcl module "litellm" { for_each = toset(["acme", "globex"]) - source = "github.com/BerriAI/litellm//terraform/litellm/aws?ref=" - tenant = each.key + source = "github.com/BerriAI/litellm//terraform/litellm/aws?ref=" + image_tag = "" + tenant = each.key env = "prod" region = "us-west-2" azs = ["us-west-2a", "us-west-2b"] @@ -485,6 +486,8 @@ provider "aws" { module "litellm" { source = "github.com/BerriAI/litellm//terraform/litellm/aws?ref=" + image_tag = "" # a git ref ships no image default, a registry version does + region = "us-west-2" tenant = "acme" env = "prod" @@ -499,12 +502,20 @@ top — set org-wide tags there, per-deployment tags via the `tags` input. ## Image pulls -The defaults pull from `ghcr.io/berriai/litellm-:v1.86.0-dev`, +The defaults pull `ghcr.io/berriai/litellm-:${image_tag}`, which is anonymous-readable. There are four images: `litellm-gateway`, `litellm-backend`, `litellm-ui`, and `litellm-migrations` (slim image used only by the one-off migration task — runs `prisma migrate deploy` against -the writer DB and exits). Bump them together when bumping LiteLLM. To pull -from a private registry: +the writer DB and exits). A published module version (Terraform Registry +`BerriAI/litellm/aws`) ships with `image_tag` stamped to its own release +tag, so bumping the module version bumps all four images with it. A git +ref of this monorepo has no published images of its own, so there +`image_tag` is required. Set `image_tag` to pin a different LiteLLM +release, or +`image_registry` to pull the same four images from a mirror. Set +`gateway_image` / `backend_image` / `ui_image` / `migrations_image` only +for a per-component override (custom build, different tag). To pull from a +private registry: - **ECR (same account)**: the execution role already has `AmazonECSTaskExecutionRolePolicy`, which grants ECR pull for repos in diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 2b235c2bad5..412e6c423f5 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -229,7 +229,7 @@ locals { gateway_metrics_container = local.metrics_enabled ? [ { name = "metrics" - image = var.gateway_image + image = local.gateway_image essential = false entryPoint = ["python", "-m", "litellm.proxy.prometheus_metrics_server"] command = ["--port", tostring(var.gateway_metrics_port)] @@ -306,7 +306,7 @@ locals { collector_container = var.collector_enabled ? [{ name = "collector" - image = var.gateway_image + image = local.gateway_image essential = false cpu = var.collector_cpu memory = var.collector_memory @@ -389,7 +389,7 @@ resource "aws_ecs_task_definition" "gateway" { merge( { name = "gateway" - image = var.gateway_image + image = local.gateway_image essential = true portMappings = [{ containerPort = 4000, protocol = "tcp" }] @@ -499,7 +499,7 @@ resource "aws_ecs_task_definition" "backend" { merge( { name = "backend" - image = var.backend_image + image = local.backend_image essential = true portMappings = [{ containerPort = 4001, protocol = "tcp" }] @@ -591,7 +591,7 @@ resource "aws_ecs_task_definition" "ui" { container_definitions = jsonencode([ { name = "ui" - image = var.ui_image + image = local.ui_image essential = true portMappings = [{ containerPort = 3000, protocol = "tcp" }] diff --git a/terraform/litellm/aws/examples/default/.terraform.lock.hcl b/terraform/litellm/aws/examples/default/.terraform.lock.hcl index 4a059b2b268..ce530fa767b 100644 --- a/terraform/litellm/aws/examples/default/.terraform.lock.hcl +++ b/terraform/litellm/aws/examples/default/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/aws" { constraints = "~> 5.60" hashes = [ "h1:Ijt7pOlB7Tr7maGQIqtsLFbl7pSMIj06TVdkoSBcYOw=", + "h1:edXOJWE4ORX8Fm+dpVpICzMZJat4AX0VRCAy/xkcOc0=", "zh:054b8dd49f0549c9a7cc27d159e45327b7b65cf404da5e5a20da154b90b8a644", "zh:0b97bf8d5e03d15d83cc40b0530a1f84b459354939ba6f135a0086c20ebbe6b2", "zh:1589a2266af699cbd5d80737a0fe02e54ec9cf2ca54e7e00ac51c7359056f274", @@ -29,6 +30,7 @@ provider "registry.terraform.io/hashicorp/random" { constraints = "~> 3.6" hashes = [ "h1:OO+IuvQJSPmWdN8AyyIEvPJbLvDQpgX/zbktoa9KsJE=", + "h1:UlBuNVuCGJ39tTv2c5gz2NRZnQbXfbIWbTzWcth5o74=", "zh:161ad0bd9a75768c82f53fb6e7172a9d8be2d4889b012645a34795031aaf1bf1", "zh:19dc9a5b17729725ccfc4f45b0500af0ee5bc6b6b160c7adb8f2bf617d2c80ea", "zh:269eda8fe42daa7974d5a34d166c3ba9defe80cde86c01e4dadcfdf2e1f05e5f", diff --git a/terraform/litellm/aws/examples/default/main.tf b/terraform/litellm/aws/examples/default/main.tf index 115003f7ddd..5e7a65f907f 100644 --- a/terraform/litellm/aws/examples/default/main.tf +++ b/terraform/litellm/aws/examples/default/main.tf @@ -38,6 +38,9 @@ module "litellm" { litellm_license = var.litellm_license ui_password = var.ui_password + image_registry = var.image_registry + image_tag = var.image_tag != "" ? var.image_tag : null + acm_certificate_arn = var.acm_certificate_arn allow_plaintext_alb = var.allow_plaintext_alb s3_force_destroy = var.s3_force_destroy diff --git a/terraform/litellm/aws/examples/default/terraform.tfvars.example b/terraform/litellm/aws/examples/default/terraform.tfvars.example index 880ecf56555..dd392d9ecc5 100644 --- a/terraform/litellm/aws/examples/default/terraform.tfvars.example +++ b/terraform/litellm/aws/examples/default/terraform.tfvars.example @@ -58,12 +58,15 @@ env = "stage" # s3_force_destroy = false # skip_final_snapshot = false -# Component images and per-task sizing/autoscaling are NOT exposed as -# variables in this example (it keeps the curated surface small). They -# default to working public GHCR images. To pin images or tune -# CPU/memory/workers/autoscaling, set those inputs directly on the -# `module "litellm"` block in main.tf — the full list is in -# ../../variables.tf — or call the module from your own root config. +# Images. image_tag is applied to all four litellm-* images. A published +# module version ships with it stamped to its own release tag, a checkout of +# the monorepo has no default and needs it set. image_registry can point at +# an ECR pull-through mirror. Per-component *_image overrides and per-task +# sizing/autoscaling are NOT exposed here. Set them directly on the +# `module "litellm"` block in main.tf (full list in ../../variables.tf) or +# call the module from your own root config. +# image_registry = "ghcr.io/berriai" +# image_tag = "" # ---------- proxy_config (mirrors helm gateway.config.proxy_config) ---------- # proxy_config = { diff --git a/terraform/litellm/aws/examples/default/variables.tf b/terraform/litellm/aws/examples/default/variables.tf index f8140266fca..27c6c95e553 100644 --- a/terraform/litellm/aws/examples/default/variables.tf +++ b/terraform/litellm/aws/examples/default/variables.tf @@ -105,6 +105,18 @@ variable "ui_password" { } # TLS — provide an ACM cert for production, or opt into HTTP-only for dev. +variable "image_registry" { + description = "Registry path prefix; images composed as /litellm-:." + type = string + default = "ghcr.io/berriai" +} + +variable "image_tag" { + description = "Tag applied to all four litellm-* images. Bump in lockstep." + type = string + default = "" +} + variable "acm_certificate_arn" { description = "ACM cert ARN for the ALB HTTPS listener. Empty → no TLS." type = string diff --git a/terraform/litellm/aws/examples/default/versions.tf b/terraform/litellm/aws/examples/default/versions.tf index 73b88e91dce..06ebe5bd2c5 100644 --- a/terraform/litellm/aws/examples/default/versions.tf +++ b/terraform/litellm/aws/examples/default/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { aws = { diff --git a/terraform/litellm/aws/locals.tf b/terraform/litellm/aws/locals.tf index b89bd486d02..c06b171d6bd 100644 --- a/terraform/litellm/aws/locals.tf +++ b/terraform/litellm/aws/locals.tf @@ -11,6 +11,11 @@ locals { # the stack can reference local.name. name = "${var.tenant}-litellm-${var.env}" + gateway_image = var.gateway_image != "" ? var.gateway_image : "${var.image_registry}/litellm-gateway:${var.image_tag}" + backend_image = var.backend_image != "" ? var.backend_image : "${var.image_registry}/litellm-backend:${var.image_tag}" + ui_image = var.ui_image != "" ? var.ui_image : "${var.image_registry}/litellm-ui:${var.image_tag}" + migrations_image = var.migrations_image != "" ? var.migrations_image : "${var.image_registry}/litellm-migrations:${var.image_tag}" + # This is a reusable module — it declares no `provider` block, so the AWS # provider's `default_tags` is the caller's concern, not ours. To keep the # same per-resource tagging the stack had when it owned the provider, the diff --git a/terraform/litellm/aws/migrations.tf b/terraform/litellm/aws/migrations.tf index e924b29eba0..57c33f69a8b 100644 --- a/terraform/litellm/aws/migrations.tf +++ b/terraform/litellm/aws/migrations.tf @@ -28,7 +28,7 @@ resource "aws_ecs_task_definition" "migrations" { container_definitions = jsonencode([{ name = "migrations" - image = var.migrations_image + image = local.migrations_image essential = true # No entryPoint/command override — the image's ENTRYPOINT runs run.py. diff --git a/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl b/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl index 5a619bc98b1..702ba37184c 100644 --- a/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl +++ b/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl @@ -14,6 +14,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/tests/collector.tftest.hcl b/terraform/litellm/aws/tests/collector.tftest.hcl index 1465130232f..082f17b2c08 100644 --- a/terraform/litellm/aws/tests/collector.tftest.hcl +++ b/terraform/litellm/aws/tests/collector.tftest.hcl @@ -13,6 +13,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" @@ -66,7 +67,7 @@ run "enabled_adds_a_sidecar_that_shares_the_gateway_transport" { assert { condition = ( - local.collector_container[0].image == var.gateway_image && + local.collector_container[0].image == local.gateway_image && local.collector_container[0].entryPoint == ["sh", "-c"] && local.collector_container[0].command == ["exec python -m litellm.proxy.collector"] && local.collector_container[0].essential == false && diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl index 7408b1179ed..c618bd11768 100644 --- a/terraform/litellm/aws/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -12,6 +12,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/tests/images.tftest.hcl b/terraform/litellm/aws/tests/images.tftest.hcl new file mode 100644 index 00000000000..2d03a46b7a3 --- /dev/null +++ b/terraform/litellm/aws/tests/images.tftest.hcl @@ -0,0 +1,104 @@ +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + image_tag = "v9.9.9" + region = "us-east-1" + tenant = "acme" + env = "test" + allow_plaintext_alb = true + azs = ["us-east-1a", "us-east-1b"] +} + +run "every_component_follows_image_registry_and_image_tag" { + command = plan + + variables { + image_registry = "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror" + image_tag = "v9.9.9" + } + + assert { + condition = ( + local.gateway_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-gateway:v9.9.9" && + local.backend_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-backend:v9.9.9" && + local.ui_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-ui:v9.9.9" && + local.migrations_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-migrations:v9.9.9" + ) + error_message = "All four images must be composed as /litellm-:." + } +} + +run "an_empty_image_tag_fails_closed_unless_every_image_is_overridden" { + command = plan + + variables { + image_tag = "" + } + + expect_failures = [var.image_tag] +} + +run "a_null_image_tag_takes_the_module_default" { + command = plan + + variables { + image_tag = null + } + + expect_failures = [var.image_tag] +} + +run "four_explicit_uris_need_no_image_tag" { + command = plan + + variables { + image_tag = "" + gateway_image = "example.com/custom/gateway:nightly" + backend_image = "example.com/custom/backend:nightly" + ui_image = "example.com/custom/ui:nightly" + migrations_image = "example.com/custom/migrations:nightly" + } + + assert { + condition = local.migrations_image == "example.com/custom/migrations:nightly" + error_message = "Explicit per-component URIs must satisfy the module without an image_tag." + } +} + +run "composed_defaults_pull_from_ghcr_at_the_image_tag" { + command = plan + + assert { + condition = alltrue([ + for image in [local.gateway_image, local.backend_image, local.ui_image, local.migrations_image] : + startswith(image, "ghcr.io/berriai/litellm-") && endswith(image, ":${var.image_tag}") + ]) + error_message = "With no overrides every component must come from ghcr.io/berriai at image_tag." + } +} + +run "a_per_component_uri_wins_over_the_composed_image" { + command = plan + + variables { + image_tag = "v9.9.9" + gateway_image = "example.com/custom/gateway:nightly" + } + + assert { + condition = ( + local.gateway_image == "example.com/custom/gateway:nightly" && + local.backend_image == "ghcr.io/berriai/litellm-backend:v9.9.9" && + local.ui_image == "ghcr.io/berriai/litellm-ui:v9.9.9" && + local.migrations_image == "ghcr.io/berriai/litellm-migrations:v9.9.9" + ) + error_message = "A per-component override must replace only that component's image." + } +} diff --git a/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl index de839d7e5be..b330d4b91c3 100644 --- a/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl +++ b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl @@ -11,6 +11,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl b/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl index 94e281faf94..ea6c343767c 100644 --- a/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl +++ b/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl @@ -11,6 +11,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 580a0cc657a..9b7a279a907 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -135,38 +135,53 @@ variable "azs" { # ---------- Component images ---------- # -# Defaults pin the four componentized images at the same release tag on -# GHCR. Override on a per-component basis in tfvars when bumping; bump them -# together when bumping the LiteLLM release. +variable "image_registry" { + description = "Registry path prefix used to compose the four LiteLLM image URIs as `/litellm-:`. Per-component overrides (`gateway_image`, `backend_image`, `ui_image`, `migrations_image`) bypass this entirely when set." + type = string + default = "ghcr.io/berriai" +} + +variable "image_tag" { + description = "Tag applied to all four litellm-* images when composed from `image_registry`. A published module version ships with this stamped to its own release tag, which matches the `v`-prefixed tag published to GHCR for the same LiteLLM release. Set it to pin a different LiteLLM release. When sourcing the module from a git ref of the monorepo it has no default and must be set." + type = string + default = "" + nullable = false + + validation { + condition = var.image_tag != "" || (var.gateway_image != "" && var.backend_image != "" && var.ui_image != "" && var.migrations_image != "") + error_message = "image_tag is empty. Published module versions (registry.terraform.io/modules/BerriAI/litellm) ship with it stamped to their own release. A git ref of the monorepo has no published images of its own, so set image_tag to a published LiteLLM release tag or set all four *_image URIs." + } +} variable "gateway_image" { - description = "Container image for the gateway (data plane, port 4000). Tag must match a tag actually published to GHCR — the split images use the `v`-prefixed semver convention." + description = "Full image URI for the gateway (data plane, port 4000). Empty (default) composes from `image_registry` + `image_tag`." type = string - default = "ghcr.io/berriai/litellm-gateway:v1.86.0-dev" + default = "" } variable "backend_image" { - description = "Container image for the backend (management API, port 4001)." + description = "Full image URI for the backend (management API, port 4001). Empty (default) composes from `image_registry` + `image_tag`." type = string - default = "ghcr.io/berriai/litellm-backend:v1.86.0-dev" + default = "" } variable "ui_image" { - description = "Container image for the UI (nginx static export, port 3000)." + description = "Full image URI for the UI (nginx static export, port 3000). Empty (default) composes from `image_registry` + `image_tag`." type = string - default = "ghcr.io/berriai/litellm-ui:v1.86.0-dev" + default = "" } variable "migrations_image" { description = <<-EOT - Container image for the one-off prisma migration task. Built from + Full image URI for the one-off prisma migration task. Empty (default) + composes from `image_registry` + `image_tag`. Built from `migrations/Dockerfile` — slim image whose ENTRYPOINT runs `python3 /app/run.py` (assembles DATABASE_URL from DATABASE_* env vars via DatabaseURLSettings, then runs `prisma migrate deploy`). Should track the same release tag as gateway/backend/ui. EOT type = string - default = "ghcr.io/berriai/litellm-migrations:v1.86.0-dev" + default = "" } # ---------- Service sizing ---------- diff --git a/terraform/litellm/aws/versions.tf b/terraform/litellm/aws/versions.tf index 73b88e91dce..06ebe5bd2c5 100644 --- a/terraform/litellm/aws/versions.tf +++ b/terraform/litellm/aws/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { aws = { diff --git a/terraform/litellm/gcp/.terraform.lock.hcl b/terraform/litellm/gcp/.terraform.lock.hcl index 013391b53d4..ad8b77ce422 100644 --- a/terraform/litellm/gcp/.terraform.lock.hcl +++ b/terraform/litellm/gcp/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/google" { constraints = "~> 6.10" hashes = [ "h1:79CwMTsp3Ud1nOl5hFS5mxQHyT0fGVye7pqpU0PPlHI=", + "h1:faTJQOetP9/RYuHwA3r2SWnuYoyzQNm4tUWZrZggcgY=", "zh:1f3513fcfcbf7ca53d667a168c5067a4dd91a4d4cccd19743e248ff31065503c", "zh:3da7db8fc2c51a77dd958ea8baaa05c29cd7f829bd8941c26e2ea9cb3aadc1e5", "zh:3e09ac3f6ca8111cbb659d38c251771829f4347ab159a12db195e211c76068bb", @@ -26,6 +27,7 @@ provider "registry.terraform.io/hashicorp/google-beta" { constraints = "~> 6.10" hashes = [ "h1:P2GiUJM1frlPtBViwKn1A9V2dVBdGuWcX80w9TdH8ZE=", + "h1:uxh4ME3hhSzVjmiWgA1IQqYqg25MV6FMVArHyA6Ki5o=", "zh:18b442bd0a05321d39dda1e9e3f1bdede4e61bc2ac62cc7a67037a3864f75101", "zh:2e387c51455862828bec923a3ec81abf63a4d998da470cf00e09003bda53d668", "zh:3942e708fa84ebe54996086f4b1398cb747fe19cbcd0be07ace528291fb35dee", @@ -45,6 +47,7 @@ provider "registry.terraform.io/hashicorp/random" { version = "3.8.1" constraints = "~> 3.6" hashes = [ + "h1:Eexl06+6J+s75uD46+WnZtpJZYRVUMB0AiuPBifK6Jc=", "h1:u8AKlWVDTH5r9YLSeswoVEjiY72Rt4/ch7U+61ZDkiQ=", "zh:08dd03b918c7b55713026037c5400c48af5b9f468f483463321bd18e17b907b4", "zh:0eee654a5542dc1d41920bbf2419032d6f0d5625b03bd81339e5b33394a3e0ae", diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index c90af26ebf2..3227dcac179 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -54,11 +54,15 @@ Then point the stack at it via `image_registry`: ```hcl image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai" -image_tag = "v1.86.0-dev" ``` -The four `litellm-:${image_tag}` URIs are composed from those -two vars. Set `gateway_image` / `backend_image` / `ui_image` / +The four `litellm-:${image_tag}` URIs are composed from +`image_registry` and `image_tag`. A published module version (Terraform +Registry `BerriAI/litellm/google`) ships with `image_tag` stamped to its +own release tag, so bumping the module version bumps all four images with +it. A git ref of this monorepo has no published images of its own, so +there `image_tag` is required. Set `image_tag` to pin a different LiteLLM +release. Set `gateway_image` / `backend_image` / `ui_image` / `migrations_image` only if you need a per-component override (custom build, different tag). @@ -485,6 +489,8 @@ provider "google-beta" { module "litellm" { source = "github.com/BerriAI/litellm//terraform/litellm/gcp?ref=" + image_tag = "" # a git ref ships no image default, a registry version does + project = "my-gcp-project" region = "us-central1" tenant = "acme" diff --git a/terraform/litellm/gcp/examples/default/.terraform.lock.hcl b/terraform/litellm/gcp/examples/default/.terraform.lock.hcl index e6285567315..55568afb834 100644 --- a/terraform/litellm/gcp/examples/default/.terraform.lock.hcl +++ b/terraform/litellm/gcp/examples/default/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/google" { constraints = "~> 6.10" hashes = [ "h1:79CwMTsp3Ud1nOl5hFS5mxQHyT0fGVye7pqpU0PPlHI=", + "h1:faTJQOetP9/RYuHwA3r2SWnuYoyzQNm4tUWZrZggcgY=", "zh:1f3513fcfcbf7ca53d667a168c5067a4dd91a4d4cccd19743e248ff31065503c", "zh:3da7db8fc2c51a77dd958ea8baaa05c29cd7f829bd8941c26e2ea9cb3aadc1e5", "zh:3e09ac3f6ca8111cbb659d38c251771829f4347ab159a12db195e211c76068bb", @@ -26,6 +27,7 @@ provider "registry.terraform.io/hashicorp/google-beta" { constraints = "~> 6.10" hashes = [ "h1:P2GiUJM1frlPtBViwKn1A9V2dVBdGuWcX80w9TdH8ZE=", + "h1:uxh4ME3hhSzVjmiWgA1IQqYqg25MV6FMVArHyA6Ki5o=", "zh:18b442bd0a05321d39dda1e9e3f1bdede4e61bc2ac62cc7a67037a3864f75101", "zh:2e387c51455862828bec923a3ec81abf63a4d998da470cf00e09003bda53d668", "zh:3942e708fa84ebe54996086f4b1398cb747fe19cbcd0be07ace528291fb35dee", @@ -46,6 +48,7 @@ provider "registry.terraform.io/hashicorp/random" { constraints = "~> 3.6" hashes = [ "h1:OO+IuvQJSPmWdN8AyyIEvPJbLvDQpgX/zbktoa9KsJE=", + "h1:UlBuNVuCGJ39tTv2c5gz2NRZnQbXfbIWbTzWcth5o74=", "zh:161ad0bd9a75768c82f53fb6e7172a9d8be2d4889b012645a34795031aaf1bf1", "zh:19dc9a5b17729725ccfc4f45b0500af0ee5bc6b6b160c7adb8f2bf617d2c80ea", "zh:269eda8fe42daa7974d5a34d166c3ba9defe80cde86c01e4dadcfdf2e1f05e5f", diff --git a/terraform/litellm/gcp/examples/default/deploystack.json b/terraform/litellm/gcp/examples/default/deploystack.json index 47d1fd914ce..732d4a187b1 100644 --- a/terraform/litellm/gcp/examples/default/deploystack.json +++ b/terraform/litellm/gcp/examples/default/deploystack.json @@ -25,7 +25,7 @@ { "name": "image_tag", "description": "Tag for the four litellm-* images (gateway, backend, ui, migrations). Bump together when bumping LiteLLM", - "default": "v1.86.0-dev" + "default": "" }, { "name": "image_registry", diff --git a/terraform/litellm/gcp/examples/default/main.tf b/terraform/litellm/gcp/examples/default/main.tf index 782d4ee4b65..9b18a6df863 100644 --- a/terraform/litellm/gcp/examples/default/main.tf +++ b/terraform/litellm/gcp/examples/default/main.tf @@ -41,7 +41,7 @@ module "litellm" { ui_password = var.ui_password image_registry = var.image_registry - image_tag = var.image_tag + image_tag = var.image_tag != "" ? var.image_tag : null lb_domains = var.lb_domains allow_plaintext_lb = var.allow_plaintext_lb diff --git a/terraform/litellm/gcp/examples/default/terraform.tfvars.example b/terraform/litellm/gcp/examples/default/terraform.tfvars.example index ec6206734e2..6456b525d08 100644 --- a/terraform/litellm/gcp/examples/default/terraform.tfvars.example +++ b/terraform/litellm/gcp/examples/default/terraform.tfvars.example @@ -38,12 +38,13 @@ env = "stage" # Images. Cloud Run rejects ghcr.io, so a real deploy must point # image_registry at an Artifact Registry remote repo (see README "Image -# pulls"); image_tag is applied to all four litellm-* images. Per-component -# *_image overrides are NOT exposed here — set them directly on the -# `module "litellm"` block in main.tf (see ../../variables.tf) if you need -# to mix-and-match versions. +# pulls"); image_tag is applied to all four litellm-* images. A published +# module version ships with it stamped to its own release tag, a checkout of +# the monorepo has no default and needs it set. Per-component *_image overrides are +# NOT exposed here. Set them directly on the `module "litellm"` block in +# main.tf (see ../../variables.tf) if you need to mix-and-match versions. # image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai" -# image_tag = "v1.86.0-dev" +# image_tag = "" # ---------- proxy_config (mirrors helm gateway.config.proxy_config) ---------- # proxy_config = { diff --git a/terraform/litellm/gcp/examples/default/variables.tf b/terraform/litellm/gcp/examples/default/variables.tf index 08b78346df3..dcf84572fd7 100644 --- a/terraform/litellm/gcp/examples/default/variables.tf +++ b/terraform/litellm/gcp/examples/default/variables.tf @@ -85,7 +85,7 @@ variable "image_registry" { variable "image_tag" { description = "Tag applied to all four litellm-* images. Bump in lockstep." type = string - default = "v1.86.0-dev" + default = "" } # TLS — provide DNS names for a managed cert, or opt into HTTP-only for dev. diff --git a/terraform/litellm/gcp/examples/default/versions.tf b/terraform/litellm/gcp/examples/default/versions.tf index a630c59afd0..55214b5c51a 100644 --- a/terraform/litellm/gcp/examples/default/versions.tf +++ b/terraform/litellm/gcp/examples/default/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { google = { diff --git a/terraform/litellm/gcp/tests/collector.tftest.hcl b/terraform/litellm/gcp/tests/collector.tftest.hcl index 7a5ea781acb..f16d02d9c7e 100644 --- a/terraform/litellm/gcp/tests/collector.tftest.hcl +++ b/terraform/litellm/gcp/tests/collector.tftest.hcl @@ -7,6 +7,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "acme-test" region = "us-central1" tenant = "acme" diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl index 999e4f0ff95..c0ad842ff95 100644 --- a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -18,6 +18,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "test-project" tenant = "tenant" env = "test" diff --git a/terraform/litellm/gcp/tests/deps_only.tftest.hcl b/terraform/litellm/gcp/tests/deps_only.tftest.hcl index 610c49d5b52..11815af9051 100644 --- a/terraform/litellm/gcp/tests/deps_only.tftest.hcl +++ b/terraform/litellm/gcp/tests/deps_only.tftest.hcl @@ -14,6 +14,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "test-project" tenant = "tenant" env = "test" @@ -112,6 +113,20 @@ run "deps_only_creates_no_runtime" { } } +run "deps_only_needs_no_image_tag" { + command = plan + + variables { + create_runtime = false + image_tag = "" + } + + assert { + condition = length(google_cloud_run_v2_service.gateway) == 0 + error_message = "deps-only mode must plan without an image tag since it deploys no images" + } +} + run "existing_network_attaches_data_stores" { command = plan diff --git a/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl index 7de62065d23..fb5446b4c56 100644 --- a/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl +++ b/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl @@ -14,6 +14,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "test-project" tenant = "tenant" env = "test" diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 412f919ab89..6161e893e02 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -144,9 +144,15 @@ variable "image_registry" { } variable "image_tag" { - description = "Tag applied to all four litellm-* images when composed from `image_registry`. Bump in lockstep when bumping LiteLLM. Must match a tag actually published to GHCR — the split images use the `v`-prefixed semver convention (e.g. `v1.86.0-dev`)." + description = "Tag applied to all four litellm-* images when composed from `image_registry`. A published module version ships with this stamped to its own release tag, which matches the `v`-prefixed tag published to GHCR for the same LiteLLM release. Set it to pin a different LiteLLM release. When sourcing the module from a git ref of the monorepo it has no default and must be set." type = string - default = "v1.86.0-dev" + default = "" + nullable = false + + validation { + condition = !var.create_runtime || var.image_tag != "" || (var.gateway_image != "" && var.backend_image != "" && var.ui_image != "" && var.migrations_image != "") + error_message = "image_tag is empty. Published module versions (registry.terraform.io/modules/BerriAI/litellm) ship with it stamped to their own release. A git ref of the monorepo has no published images of its own, so set image_tag to a published LiteLLM release tag or set all four *_image URIs." + } } variable "gateway_image" { diff --git a/terraform/litellm/gcp/versions.tf b/terraform/litellm/gcp/versions.tf index a630c59afd0..55214b5c51a 100644 --- a/terraform/litellm/gcp/versions.tf +++ b/terraform/litellm/gcp/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { google = { diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 3ad80af10f0..f50096add19 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -78,6 +78,7 @@ - {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"} - {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"} - {id: mgmt.fallback_management.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "fallback_management_endpoints.py", rationale: "Fallback config (smoke)"} +- {id: mgmt.fallback_management.update.config_owned_refused, module: mgmt, tier: P1, surface: api, assertions: [config_owned_refused], source: "fallback_management_endpoints.py", rationale: "A /fallback write to a fallbacks key the config file sets is refused and the config rule keeps routing"} - {id: mgmt.workflow.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "workflow_management_endpoints.py", rationale: "Workflow tracking (smoke)"} - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9198f6c54d6..90a9e9367eb 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -499,6 +499,29 @@ def get_external[R: BaseModel]( return classify(resp, response_type) +def get_text_external( + url: str, + *, + headers: BaseModel | None = None, + timeout: float = 30.0, + ca_bundle: str | None = None, +) -> str | NetworkError: + """GET an absolute URL outside the proxy whose body is plain text (e.g. a pod log), not JSON. + A non-2xx status is reported as a NetworkError so the caller can fail loudly.""" + try: + resp = requests.get( + url, + headers={"Accept": "text/plain", **(_headers(headers) if headers is not None else {})}, + timeout=timeout, + verify=ca_bundle if ca_bundle is not None else True, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + if not (200 <= resp.status_code < 300): + return NetworkError(message=f"{resp.status_code} from {url}: {resp.text[:300]}") + return resp.text + + def post_external[R: BaseModel]( url: str, *, diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index 12f707cf6f1..7fd51bd6e8c 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -11,15 +11,17 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix - Vertex (gemini-2.5-flash): explicit context caching via ``cache_control`` with a 5-minute ttl. litellm builds the Vertex cache before the generate call, so a never-seen prefix must come back cached on its very first call - (Gemini's implicit caching cannot hit a cold prefix), the cached count must - cover the marked block, and the spend row must be billed below the uncached - price of the cached tokens. Vertex's cache create is nondeterministic: + (Gemini's implicit caching cannot hit a cold prefix). That call bills the + cache write as ``cache_creation_input_tokens`` on top of Vertex's own + prompt count, the cached count must cover the marked block, the spend row + must carry the write at the input rate and the read below it, and an + identical second call must read the cache without building a new one. + Vertex's cache create is nondeterministic: identical bodies come back 200 or with the minimum-token 400 ("The cached content is of 1 tokens"), in failure bursts of 45 seconds and more, so up to eight never-seen prefixes are tried with a pause after each rejection. The - billing check prices the cached tokens rather than prompt_tokens, which - Vertex reports inclusive of the cached prefix on some calls and exclusive - of it on others. + checks hold whether Vertex reports its prompt count inclusive of the cached + prefix or exclusive of it, which it has done both ways. - Anthropic (claude-haiku-4-5, direct): the same ``cache_control`` prefix over the OpenAI-compatible route; the second call must report cache-read tokens > 0. - OpenAI (gpt-5.6): automatic prompt caching needs no request marker, so the @@ -42,6 +44,8 @@ from __future__ import annotations import time from collections.abc import Callable +from dataclasses import dataclass +from functools import partial from typing import Final import pytest @@ -60,6 +64,7 @@ from models import ( ChatTool, ChatToolFunction, ChatToolResultTurn, + CustomPricing, LiteLLMParamsBody, PromptCacheOptions, RichMessage, @@ -203,19 +208,26 @@ def _assert_openai_writes_then_reads_cache(client: PassthroughClient, key: str, assert read_tokens > 0, f"{body.model}: the identical second call read no cache ({second.usage})" -def _cold_cache_call(send: Callable[[str], Result[ChatResponse]]) -> ChatResponse | None: - result: Final = send(_cacheable_prefix()) +@dataclass(frozen=True, slots=True) +class EngagedColdCall: + prefix: str + completion: ChatResponse + + +def _cold_cache_call(send: Callable[[str], Result[ChatResponse]]) -> EngagedColdCall | None: + prefix: Final = _cacheable_prefix() + result: Final = send(prefix) match result: case UnknownApiError(status_code=400, body=body) if VERTEX_CACHE_REJECTION_MARKER in body: return None case _: - return unwrap(result) + return EngagedColdCall(prefix=prefix, completion=unwrap(result)) -def _first_engaged_cold_call(send: Callable[[str], Result[ChatResponse]]) -> ChatResponse | None: +def _first_engaged_cold_call(send: Callable[[str], Result[ChatResponse]]) -> EngagedColdCall | None: for attempt in range(1, VERTEX_COLD_CALL_ATTEMPTS + 1): candidate = _cold_cache_call(send) - if candidate is not None and _cached_read_tokens(candidate.usage) >= VERTEX_MINIMUM_CACHED_TOKENS: + if candidate is not None and _cached_read_tokens(candidate.completion.usage) >= VERTEX_MINIMUM_CACHED_TOKENS: return candidate if attempt < VERTEX_COLD_CALL_ATTEMPTS: print( @@ -227,32 +239,61 @@ def _first_engaged_cold_call(send: Callable[[str], Result[ChatResponse]]) -> Cha return None -def _first_cold_call_reads_cache(model: str, send: Callable[[str], Result[ChatResponse]]) -> ChatResponse: - completion: Final = _first_engaged_cold_call(send) - assert completion is not None, ( +def _cache_creation_tokens(usage: Usage | None) -> int: + if usage is None: + return 0 + return usage.cache_creation_input_tokens or 0 + + +def _first_cold_call_creates_and_reads_cache( + model: str, send: Callable[[str], Result[ChatResponse]] +) -> EngagedColdCall: + engaged: Final = _first_engaged_cold_call(send) + assert engaged is not None, ( f"{model}: {VERTEX_COLD_CALL_ATTEMPTS} never-seen prompts marked with cache_control, spread over " f"{VERTEX_COLD_CALL_PAUSE_SECONDS * (VERTEX_COLD_CALL_ATTEMPTS - 1):.0f}s, were each either rejected by " f"Vertex's minimum-token check or served with fewer than {VERTEX_MINIMUM_CACHED_TOKENS} cached tokens on " "their first call; explicit context caching did not engage" ) + completion: Final = engaged.completion assert completion.choices, f"{model}: cached call returned no choices: {completion}" usage: Final = completion.usage - cached: Final = _cached_read_tokens(usage) assert usage and usage.prompt_tokens, f"{model}: cached completion carried no prompt_tokens: {usage}" - assert cached >= CACHED_SHARE_OF_PROMPT * usage.prompt_tokens, ( - f"{model}: only {cached} of {usage.prompt_tokens} prompt tokens were served from the cache; the " - "cache_control block was not cached whole" + created: Final = _cache_creation_tokens(usage) + assert created > 0, ( + f"{model}: the call that built the Vertex cache reported no cache_creation_input_tokens ({usage}); " + "the cache write went unbilled" ) - return completion + generate_prompt_tokens: Final = usage.prompt_tokens - created + assert generate_prompt_tokens > 0, ( + f"{model}: prompt_tokens {usage.prompt_tokens} does not cover the {created} cache-creation tokens plus " + "the generate call's own prompt" + ) + cached: Final = _cached_read_tokens(usage) + assert cached >= CACHED_SHARE_OF_PROMPT * generate_prompt_tokens, ( + f"{model}: only {cached} of the generate call's {generate_prompt_tokens} prompt tokens were served from " + "the cache; the cache_control block was not cached whole" + ) + return engaged -def _input_rate(client: PassthroughClient, model: str) -> float: +def _assert_second_call_reuses_cache(model: str, send: Callable[[str], Result[ChatResponse]], prefix: str) -> None: + usage: Final = unwrap(send(prefix)).usage + assert _cached_read_tokens(usage) > 0, f"{model}: the identical second call read no cache ({usage})" + assert _cache_creation_tokens(usage) == 0, ( + f"{model}: the identical second call built and billed a new cache ({usage}) instead of reusing the first one" + ) + + +def _pricing(client: PassthroughClient, model: str) -> CustomPricing: entry: Final = next((row for row in client.proxy.model_info() if row.model_name == model), None) assert entry and entry.model_info.input_cost_per_token, f"/model/info resolved no input rate for {model}" - return entry.model_info.input_cost_per_token + return entry.model_info -def _assert_billed_below_uncached_prompt(client: PassthroughClient, model: str, completion: ChatResponse) -> None: +def _assert_billed_for_creation_and_discounted_read( + client: PassthroughClient, model: str, completion: ChatResponse +) -> None: assert completion.id, f"{model}: cached completion carried no id to find its spend row by" usage: Final = completion.usage assert usage and usage.prompt_tokens, f"{model}: cached completion carried no prompt_tokens: {usage}" @@ -262,11 +303,30 @@ def _assert_billed_below_uncached_prompt(client: PassthroughClient, model: str, assert row.prompt_tokens == usage.prompt_tokens, ( f"{model}: spend row prompt_tokens {row.prompt_tokens} != response prompt_tokens {usage.prompt_tokens}" ) + assert row.spend is not None, f"{model}: spend row carried no spend: {row}" + pricing: Final = _pricing(client, model) + input_rate: Final = pricing.input_cost_per_token or 0 + created: Final = _cache_creation_tokens(usage) cached: Final = _cached_read_tokens(usage) - uncached_read_cost: Final = cached * _input_rate(client, model) - assert row.spend is not None and row.spend < uncached_read_cost, ( - f"{model}: spend {row.spend} is not below the uncached price of the {cached} tokens read from the cache " - f"({uncached_read_cost}); cache-read pricing was not applied" + # Google bills cache creation at the standard input price (read 2026-10-16): + # https://docs.cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-overview + creation_cost: Final = created * input_rate + assert row.spend >= creation_cost, ( + f"{model}: spend {row.spend} is below the {creation_cost} the {created} cache-creation tokens cost at the " + "input rate; the cache write was not billed" + ) + read_rate: Final = pricing.cache_read_input_token_cost + assert read_rate is not None and read_rate < input_rate, ( + f"/model/info advertises no discounted cache-read rate for {model}: {pricing}" + ) + undiscounted_cost: Final = pricing.token_cost( + created + max(usage.prompt_tokens - created, cached), usage.completion_tokens or 0 + ) + half_the_read_discount: Final = cached * (input_rate - read_rate) / 2 + assert row.spend < undiscounted_cost - half_the_read_discount, ( + f"{model}: spend {row.spend} saved less than half the cache-read discount on its {cached} cached tokens " + f"against {undiscounted_cost}, the price of creating the cache and reading it uncached; cache-read pricing " + "was not applied" ) @@ -326,10 +386,10 @@ class TestCacheControl: ) resources.defer(lambda: client.proxy.delete_model(model_id)) key = resources.key() - completion: Final = _first_cold_call_reads_cache( - model, lambda prefix: _cache_chat(client, key, model, prefix, ttl=VERTEX_CACHE_TTL) - ) - _assert_billed_below_uncached_prompt(client, model, completion) + send: Final = partial(_cache_chat, client, key, model, ttl=VERTEX_CACHE_TTL) + engaged: Final = _first_cold_call_creates_and_reads_cache(model, send) + _assert_billed_for_creation_and_discounted_read(client, model, engaged.completion) + _assert_second_call_reuses_cache(model, send, engaged.prefix) @pytest.mark.covers( "llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works", diff --git a/tests/e2e/management/test_config_misc_endpoints_e2e.py b/tests/e2e/management/test_config_misc_endpoints_e2e.py index b7ab311cc18..bd9192fa0f0 100644 --- a/tests/e2e/management/test_config_misc_endpoints_e2e.py +++ b/tests/e2e/management/test_config_misc_endpoints_e2e.py @@ -31,7 +31,7 @@ import pytest from pydantic import BaseModel from e2e_config import unique_marker -from e2e_http import NoBody, Success, unwrap, unwrap_status +from e2e_http import NoBody, Result, Success, UnknownApiError, unwrap, unwrap_status from e2e_metadata import Domain, Route, Subject, meta from lifecycle import ResourceManager from management_client import ManagementClient @@ -141,12 +141,15 @@ class ComplianceResponse(BaseModel): # ---- fallback management --------------------------------------------------- -class FallbackShape(BaseModel): +class FallbackTarget(BaseModel): model: str - fallback_models: list[str] fallback_type: str +class FallbackShape(FallbackTarget): + fallback_models: list[str] + + class FallbackCreateBody(FallbackShape): pass @@ -163,6 +166,28 @@ class FallbackGetResponse(FallbackShape): pass +class FallbackDeleteResponse(FallbackTarget): + message: str + + +class RouterFallbackValues(BaseModel): + fallbacks: list[dict[str, list[str]]] | None = None + + +class RouterFallbackOwnership(BaseModel): + current_values: RouterFallbackValues + source: dict[str, str] + + +class ConfigOwnedDetail(BaseModel): + keys: list[str] + section: str + + +class ConfigOwnedRejection(BaseModel): + detail: ConfigOwnedDetail + + # ---- jwt key mapping ------------------------------------------------------- @@ -365,20 +390,19 @@ class TestFallbackManagement: @pytest.mark.covers("mgmt.fallback_management.update.happy_path") @meta(Subject(domain=Domain.ROUTING)) def test_create_persists_and_is_read_back(self, client: ManagementClient, resources: ResourceManager) -> None: - primary = f"e2e-fallback-primary-{unique_marker()}" - secondary = f"e2e-fallback-secondary-{unique_marker()}" - params = LiteLLMParamsBody(model="openai/gpt-5.5", api_key="e2e-dummy-key") - primary_id = client.proxy.create_model(primary, params) - resources.defer(lambda: client.proxy.delete_model(primary_id)) - secondary_id = client.proxy.create_model(secondary, params) - resources.defer(lambda: client.proxy.delete_model(secondary_id)) - resources.defer(lambda: self._delete_fallback(client, primary)) + ownership = self._router_fallback_ownership(client) + assert ownership.source.get("context_window_fallbacks") != "config", ( + "the proxy config sets router_settings.context_window_fallbacks, so /fallback cannot write it and this " + "round trip needs a fallback type the database owns" + ) + primary, secondary = self._register_fallback_pair(client, resources) + resources.defer(lambda: self._delete_fallback(client, primary, "context_window")) created = unwrap( client.proxy.transport.post( "/fallback", headers=client.proxy.transport.master, - json=FallbackCreateBody(model=primary, fallback_models=[secondary], fallback_type="general"), + json=FallbackCreateBody(model=primary, fallback_models=[secondary], fallback_type="context_window"), response_type=FallbackResponse, ) ) @@ -388,13 +412,7 @@ class TestFallbackManagement: ) def read_back() -> FallbackGetResponse | None: - result = client.proxy.transport.get( - f"/fallback/{primary}", - headers=client.proxy.transport.master, - params=FallbackGetParams(fallback_type="general"), - response_type=FallbackGetResponse, - ) - match result: + match self._get_fallback(client, primary, "context_window"): case Success(data=data) if secondary in data.fallback_models: return data case _: @@ -405,14 +423,105 @@ class TestFallbackManagement: f"GET /fallback/{primary} reports fallbacks {got.fallback_models}, configured [{secondary!r}]" ) + deleted = unwrap(self._delete_fallback(client, primary, "context_window")) + assert deleted.model == primary and deleted.fallback_type == "context_window", ( + f"DELETE /fallback/{primary} answered for model={deleted.model!r} type={deleted.fallback_type!r}" + ) + + def gone() -> bool | None: + match self._get_fallback(client, primary, "context_window"): + case UnknownApiError(status_code=404): + return True + case _: + return None + + _ = _poll(client, gone, f"GET /fallback/{primary} still reports a fallback after DELETE /fallback/{primary}") + + @pytest.mark.covers("mgmt.fallback_management.update.config_owned_refused") + @meta(Subject(domain=Domain.ROUTING)) + def test_write_to_config_owned_fallbacks_is_refused( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + before = self._router_fallback_ownership(client) + config_rules = before.current_values.fallbacks + assert before.source.get("fallbacks") == "config" and config_rules, ( + "this test needs router_settings.fallbacks set in the proxy config, but /router/settings reports " + f"source={before.source.get('fallbacks')!r} value={config_rules}" + ) + primary, secondary = self._register_fallback_pair(client, resources) + resources.defer(lambda: self._delete_fallback(client, primary, "general")) + + result = client.proxy.transport.post( + "/fallback", + headers=client.proxy.transport.master, + json=FallbackCreateBody(model=primary, fallback_models=[secondary], fallback_type="general"), + response_type=FallbackResponse, + ) + match result: + case UnknownApiError(status_code=400, body=body): + rejection = ConfigOwnedRejection.model_validate_json(body).detail + case _: + pytest.fail(f"/fallback accepted a write to the config-owned router_settings.fallbacks: {result}") + assert rejection.keys == ["fallbacks"] and rejection.section == "router_settings", ( + f"/fallback refused the write naming keys={rejection.keys} section={rejection.section!r}, " + "expected router_settings.fallbacks" + ) + + after = self._router_fallback_ownership(client).current_values.fallbacks + assert after == config_rules, ( + f"the refused write changed router_settings.fallbacks from {config_rules} to {after}" + ) + match self._get_fallback(client, primary, "general"): + case UnknownApiError(status_code=404): + pass + case refused_rule: + pytest.fail(f"the refused rule for {primary} is served anyway: {refused_rule}") + config_model, config_targets = next(iter(config_rules[0].items())) + served = unwrap(self._get_fallback(client, config_model, "general")) + assert served.fallback_models == config_targets, ( + f"GET /fallback/{config_model} reports {served.fallback_models}, " + f"the config file routes it to {config_targets}" + ) + @staticmethod - def _delete_fallback(client: ManagementClient, model: str) -> None: - _ = client.proxy.transport.delete( + def _register_fallback_pair(client: ManagementClient, resources: ResourceManager) -> tuple[str, str]: + primary = f"e2e-fallback-primary-{unique_marker()}" + secondary = f"e2e-fallback-secondary-{unique_marker()}" + params = LiteLLMParamsBody(model="openai/gpt-5.5", api_key="e2e-dummy-key") + primary_id = client.proxy.create_model(primary, params) + resources.defer(lambda: client.proxy.delete_model(primary_id)) + secondary_id = client.proxy.create_model(secondary, params) + resources.defer(lambda: client.proxy.delete_model(secondary_id)) + return primary, secondary + + @staticmethod + def _router_fallback_ownership(client: ManagementClient) -> RouterFallbackOwnership: + return unwrap( + client.proxy.transport.get( + "/router/settings", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=RouterFallbackOwnership, + ) + ) + + @staticmethod + def _get_fallback(client: ManagementClient, model: str, fallback_type: str) -> Result[FallbackGetResponse]: + return client.proxy.transport.get( + f"/fallback/{model}", + headers=client.proxy.transport.master, + params=FallbackGetParams(fallback_type=fallback_type), + response_type=FallbackGetResponse, + ) + + @staticmethod + def _delete_fallback(client: ManagementClient, model: str, fallback_type: str) -> Result[FallbackDeleteResponse]: + return client.proxy.transport.delete( f"/fallback/{model}", headers=client.proxy.transport.master, json=NoBody(), - params=FallbackGetParams(fallback_type="general"), - response_type=NoBody, + params=FallbackGetParams(fallback_type=fallback_type), + response_type=FallbackDeleteResponse, ) diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py new file mode 100644 index 00000000000..0d67a98257f --- /dev/null +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -0,0 +1,49 @@ +"""Where the gateway holds its event loop for 100 ms or more: the mixed case's callers, model groups, guardrails +and body mix (short chats, 300k to 1M-token contexts, tool calls, half streamed) at 1 rps for 120 requests against +a gateway running the stdlib asyncio loop in debug mode, which logs every loop step over 0.1 s with the task that +ran it. No latency, CPU or datastore ceilings: the one check is that no step over 0.1 s comes from a location that +is not on the allowlist below, and that the listed ones do not show up more often than they do today. A stream +that stops yielding (the 46925 fix), a provider handler that re-walks a 1M-token body on the loop (46967), or a +new hook that deep-copies the request inline all land here as a new row. Runs on its own stack: the gateway needs +`PYTHONASYNCIODEBUG=1` and `LITELLM_EVENT_LOOP=asyncio` (uvloop reports nothing and crashed under debug locally), +and the runner `E2E_PERF_EVENT_LOOP=1`, so the nightly build, which sets neither, skips it. The allowlist counts +are twice the highest count any of ten back-to-back local runs showed on that row (rows that showed nothing in +those ten but did in the earlier before/after legs for 46925 and 46967 get 2), so each run sits at or under half +of every cap and a listed location doubling in frequency is a failure, not a flake. The +`garbage collection` bucket is the gen-2 passes the collector logs in debug mode (0.6 to 0.7 s each, one every +minute or two of this traffic, plus the ramp while the heap fills on a local whole-file run), which asyncio would +otherwise charge to a random task; its count is the signal for a change that makes much more garbage.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Final + +from setup.event_loop import EventLoopStalls, KnownStall + +from cases.mixed_size_requests_2rps import CASE as MIXED + +CASE: Final = replace( + MIXED, + name="event_loop_stalls", + traffic=replace(MIXED.traffic, rps=1.0, total_requests=120, concurrency=10), + limits=(), + event_loop=EventLoopStalls( + allowlist=( + KnownStall("garbage collection", max_count=4), + KnownStall("LifespanOn.main", max_count=4), + KnownStall("RequestResponseCycle.run_asgi", max_count=2), + KnownStall("with_proxy_compaction_executor", max_count=2), + KnownStall("Logging.async_success_handler", max_count=2), + KnownStall("Logging.dispatch_success_handlers", max_count=2), + KnownStall("_TrackedPrismaEngine", max_count=2), + KnownStall("prisma/actions.py", max_count=2), + KnownStall("SpendCounterReseed", max_count=2), + KnownStall("Subscriber._run", max_count=2), + KnownStall("RedisBatch.flush", max_count=2), + KnownStall("run_coroutine_job", max_count=2), + KnownStall("user_api_key_auth.py", max_count=2), + KnownStall("route_streaming_logging_to_handler", max_count=2), + ), + ), +) diff --git a/tests/e2e/perf/setup/event_loop.py b/tests/e2e/perf/setup/event_loop.py new file mode 100644 index 00000000000..dee45c374d9 --- /dev/null +++ b/tests/e2e/perf/setup/event_loop.py @@ -0,0 +1,246 @@ +"""Event-loop holds, read from the gateway's own asyncio debug output. A gateway started with +`PYTHONASYNCIODEBUG=1` on the stdlib loop (`LITELLM_EVENT_LOOP=asyncio`, uvloop never reports them) logs every +callback or task step over `slow_callback_duration` (0.1 s) as +`Executing ...> took X seconds`. A case with `event_loop` set reads +the gateway log for its traffic window, groups those lines by coroutine and location, and fails on a location that +is not on its allowlist of known holds or shows up more often than the allowlist gives it. asyncio names the task, +not the frame that held the loop, so every request handler lands under `RequestResponseCycle.run_asgi`; the +allowlist caps how often a known location may appear instead of ignoring it. The log comes from the gateway pods +(`pods/log` with the runner's service account, pods picked by `E2E_PERF_GATEWAY_POD_SELECTOR`) or, on a local run, +from the whole file named by `E2E_PERF_GATEWAY_LOG_FILE` (no time window, so startup holds are in it too). +A garbage collection pass pauses the loop too, and asyncio charges it to whatever task was running, so the gateway +logs every pass over the same threshold (`Garbage collection gen=2 collected=N took X seconds`, debug mode only); +a hold that such a pass accounts for is reported under `garbage collection` instead of the task it landed on.""" + +from __future__ import annotations + +import os +import re +from collections.abc import Iterable +from dataclasses import dataclass, replace +from datetime import datetime +from itertools import groupby +from pathlib import Path +from typing import Final +from urllib.parse import quote + +from e2e_http import AuthHeaders, NetworkError, Success, get_external, get_text_external +from pydantic import BaseModel + +from setup.podcpu import SELECTOR_ENV, in_cluster + +OPT_IN_ENV: Final = "E2E_PERF_EVENT_LOOP" +LOG_FILE_ENV: Final = "E2E_PERF_GATEWAY_LOG_FILE" +_API: Final = "https://kubernetes.default.svc/api/v1/namespaces" +_LINE: Final = re.compile(r"Executing <(?P.*?)> took (?P[0-9.]+) seconds") +_CORO: Final = re.compile(r"coro=<(?P[^ ]+)") +_AT: Final = re.compile(r" at (?P[^ >]+:\d+)") +_GC_LINE: Final = re.compile( + r"Garbage collection gen=(?P\d+) collected=\d+ took (?P[0-9.]+) seconds" +) +GC_BUCKET: Final = "garbage collection" + + +@dataclass(frozen=True, slots=True) +class KnownStall: + """A hold the gateway has today: `where` is a substring of the coroutine name or its file:line, `max_count` the + most times it may show up in one run before the case fails.""" + + where: str + max_count: int + + +@dataclass(frozen=True, slots=True) +class EventLoopStalls: + allowlist: tuple[KnownStall, ...] + max_hold_seconds: float = 0.1 + + +@dataclass(frozen=True, slots=True) +class GcPass: + generation: int + seconds: float + + +@dataclass(frozen=True, slots=True) +class Stall: + coroutine: str + where: str + seconds: float + gc_seconds: float = 0.0 + gc_generation: int | None = None + + +@dataclass(frozen=True, slots=True) +class StallRow: + coroutine: str + where: str + count: int + total_seconds: float + max_seconds: float + allowed_by: KnownStall | None + + @property + def offending(self) -> bool: + return self.allowed_by is None or self.count > self.allowed_by.max_count + + +def _short(path: str) -> str: + if "site-packages/" in path: + return path.rsplit("site-packages/", 1)[1] + if "/litellm/" in path: + return "litellm/" + path.rsplit("/litellm/", 1)[1] + return path + + +def _parse_line(line: str) -> Stall | None: + match: Final = _LINE.search(line) + if match is None: + return None + desc: Final = match.group("desc") + coro: Final = _CORO.search(desc) + at: Final = _AT.search(desc) + return Stall( + coroutine=coro.group("name") if coro else desc[:80], + where=_short(at.group("where")) if at else "", + seconds=float(match.group("seconds")), + ) + + +def _parse_event(line: str) -> Stall | GcPass | None: + match: Final = _GC_LINE.search(line) + if match is None: + return _parse_line(line) + return GcPass(generation=int(match.group("generation")), seconds=float(match.group("seconds"))) + + +def _with_passes(stall: Stall, passes: tuple[GcPass, ...]) -> Stall: + fitting: Final = tuple(gc_pass for gc_pass in passes if gc_pass.seconds <= stall.seconds) + if not fitting: + return stall + longest: Final = max(fitting, key=lambda gc_pass: gc_pass.seconds) + return replace(stall, gc_seconds=longest.seconds, gc_generation=longest.generation) + + +def parse_stalls(lines: Iterable[str]) -> tuple[Stall, ...]: + """Holds in log order, each carrying the collector passes logged since the previous hold: a pass runs inside the + task asyncio then blames and is only logged when it is itself over the threshold, so its hold line follows it.""" + events: Final = tuple(event for event in map(_parse_event, lines) if event is not None) + holds: Final = tuple((index, event) for index, event in enumerate(events) if isinstance(event, Stall)) + return tuple( + _with_passes(stall, tuple(event for event in events[previous + 1 : index] if isinstance(event, GcPass))) + for previous, (index, stall) in zip((-1, *(index for index, _ in holds)), holds) + ) + + +def _reported(spec: EventLoopStalls, stall: Stall) -> Stall: + """A hold the collector accounts for, what is left under the threshold, is the collector's, not the task's.""" + if stall.gc_generation is None: + return stall + if stall.seconds - stall.gc_seconds >= spec.max_hold_seconds: + return stall + return Stall(coroutine=GC_BUCKET, where=f"gen={stall.gc_generation}", seconds=stall.seconds) + + +def _allowed_by(spec: EventLoopStalls, coroutine: str, where: str) -> KnownStall | None: + return next((known for known in spec.allowlist if known.where in coroutine or known.where in where), None) + + +def summarize(spec: EventLoopStalls, stalls: Iterable[Stall]) -> tuple[StallRow, ...]: + """One row per (coroutine, location) over the hold threshold, worst total first.""" + reported: Final = tuple(_reported(spec, stall) for stall in stalls) + over: Final = sorted( + (stall for stall in reported if stall.seconds >= spec.max_hold_seconds), key=lambda s: (s.coroutine, s.where) + ) + rows: Final = tuple( + StallRow( + coroutine=coroutine, + where=where, + count=len(group), + total_seconds=sum(stall.seconds for stall in group), + max_seconds=max(stall.seconds for stall in group), + allowed_by=_allowed_by(spec, coroutine, where), + ) + for (coroutine, where), group in ( + (key, tuple(items)) for key, items in groupby(over, key=lambda s: (s.coroutine, s.where)) + ) + ) + return tuple(sorted(rows, key=lambda row: -row.total_seconds)) + + +def offenders(rows: Iterable[StallRow]) -> tuple[str, ...]: + return tuple( + f"{row.coroutine} {row.where}: {row.count}x, max {row.max_seconds:.2f} s" + + ( + f" (allowlist {row.allowed_by.where} permits {row.allowed_by.max_count})" + if row.allowed_by + else " (not on the allowlist)" + ) + for row in rows + if row.offending + ) + + +def render_stalls(spec: EventLoopStalls, rows: Iterable[StallRow]) -> str: + header: Final = ( + f" event-loop holds >= {spec.max_hold_seconds:.2f} s in the gateway log (count total_s max_s where)" + ) + body: Final = tuple( + f" {'FAIL' if row.offending else 'ok '} {row.count:5d} {row.total_seconds:8.2f} {row.max_seconds:6.2f} " + f"{row.coroutine} {row.where}" + for row in rows + ) + return "\n".join((header, *body)) if body else header + "\n none" + + +class _PodName(BaseModel): + name: str + + +class _Pod(BaseModel): + metadata: _PodName + + +class _PodList(BaseModel): + items: list[_Pod] + + +def _pod_log_lines(url: str, headers: AuthHeaders, ca_bundle: str) -> tuple[str, ...]: + text: Final = get_text_external(url, headers=headers, timeout=60.0, ca_bundle=ca_bundle) + if isinstance(text, NetworkError): + raise RuntimeError(f"gateway pod log unreadable: {text.message}") + return tuple(text.splitlines()) + + +def _pod_logs(since: datetime) -> tuple[str, ...] | None: + cluster: Final = in_cluster() + selector: Final = os.environ.get(SELECTOR_ENV, "") + if cluster is None or not selector: + return None + token, ca_bundle, namespace = cluster + headers: Final = AuthHeaders(authorization=f"Bearer {token}") + pods: Final = get_external( + f"{_API}/{namespace}/pods?labelSelector={quote(selector)}", + response_type=_PodList, + headers=headers, + timeout=60.0, + ca_bundle=ca_bundle, + ) + if not isinstance(pods, Success): + raise RuntimeError(f"gateway pod list unreadable: {pods}") + since_param: Final = quote(since.strftime("%Y-%m-%dT%H:%M:%SZ")) + return tuple( + line + for pod in pods.data.items + for line in _pod_log_lines( + f"{_API}/{namespace}/pods/{pod.metadata.name}/log?sinceTime={since_param}", headers, ca_bundle + ) + ) + + +def read_gateway_log(since: datetime) -> tuple[str, ...] | None: + """Gateway log lines for the traffic window, or None when neither a pod nor a local file is reachable.""" + local: Final = os.environ.get(LOG_FILE_ENV, "") + if local: + return tuple(Path(local).read_text(errors="replace").splitlines()) + return _pod_logs(since) diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py index 1609617bd73..ccc7cf698f8 100644 --- a/tests/e2e/perf/setup/locustfile.py +++ b/tests/e2e/perf/setup/locustfile.py @@ -19,9 +19,10 @@ from typing import Final, cast import gevent from locust import FastHttpUser, task from locust.contrib.fasthttp import FastResponse, ResponseContextManager -from setup.wire import Endpoint, WireRecord, WireRequest, request_body from pydantic import TypeAdapter +from setup.wire import Endpoint, PlanProgress, WireRecord, WireRequest, request_body, response_id + _PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest]) _JSON: Final = TypeAdapter[object](object) _PLAN: Final = _PLAN_ADAPTER.validate_json(Path(os.environ["PERF_PLAN_FILE"]).read_bytes()) @@ -33,7 +34,8 @@ _INTERVAL_S: Final = _USERS / _RPS _JITTER: Final = 0.25 _RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan") _RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8") -_written = 0 +_PROGRESS: Final = PlanProgress(len(_PLAN)) +_STOP_FILE: Final = os.environ.get("PERF_STOP_FILE") BLOCK_REASON: Final = "perf-test-block" _TOKEN_MAX_AGE_S: Final = 240.0 _tokens: dict[str, tuple[str, float]] = {} @@ -155,10 +157,9 @@ def _judge(item: WireRequest, status: int, body_text: str) -> tuple[bool, str | def _record(record: WireRecord) -> None: - global _written _ = _RECORDS.write(record.model_dump_json() + "\n") _RECORDS.flush() - _written += 1 + _PROGRESS.completed += 1 def _authorization(user: FastHttpUser, item: WireRequest) -> str: @@ -203,7 +204,7 @@ class PlanUser(FastHttpUser): a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate.""" - network_timeout = 120.0 + network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120")) connection_timeout = 30.0 def on_start(self) -> None: @@ -221,9 +222,10 @@ class PlanUser(FastHttpUser): @task def send(self) -> None: - item: Final = next(_NEXT, None) + take: Final = _PROGRESS.take(stop=_STOP_FILE is not None and Path(_STOP_FILE).exists()) + item: Final = next(_NEXT, None) if take else None if item is None: - if _written >= len(_PLAN): + if _PROGRESS.drained: _quit(self) return started_at: Final = time.time() @@ -237,6 +239,7 @@ class PlanUser(FastHttpUser): error: str | None = None header_s = 0.0 response: FastResponse | None = None + identifier: str | None = None # rebind-ok: the ID arrives in a JSON body or the first SSE event try: headers: Final = {"Authorization": _authorization(self, item), "Content-Type": "application/json"} context: ResponseContextManager @@ -265,6 +268,7 @@ class PlanUser(FastHttpUser): buffer += line events, buffer = _events(buffer) for event in events: + identifier = identifier or response_id(event) if first_token_s is None and token_in_event(item.endpoint, event): first_token_s = now text += delta_of_event(item.endpoint, event) @@ -272,7 +276,9 @@ class PlanUser(FastHttpUser): body_text = context.text or "" first_byte_s = time.perf_counter() - started if status == 200: - text = text_of_json(item.endpoint, _as_dict(_JSON.validate_json(body_text))) + payload: Final = _as_dict(_JSON.validate_json(body_text)) + identifier = response_id(payload) + text = text_of_json(item.endpoint, payload) ok, error = _judge(item, status, body_text) if ok: context.success() @@ -302,7 +308,8 @@ class PlanUser(FastHttpUser): mock_delay_s=float(delay_header) if delay_header else None, call_id=_header(response, "x-litellm-call-id") if response is not None else None, error=error, + response_id=identifier, ) ) - if _written >= len(_PLAN): + if _PROGRESS.drained: _quit(self) diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index 37ad640db62..2879866f85b 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -13,6 +13,8 @@ from collections.abc import Mapping from dataclasses import dataclass, field from typing import Final, Literal +from setup.event_loop import EventLoopStalls + Endpoint = Literal["chat_completions", "responses", "messages"] MockProvider = Literal["openai", "anthropic", "bedrock"] GuardrailMode = Literal["pre_call", "post_call"] @@ -153,6 +155,7 @@ class PerformanceTestCase: background_cap_exempt_tables: frozenset[str] = frozenset() """Tables whose background write is one upsert per writer flush and, at this case's request rate, lands above the `postgres_max_background_statement_per_request` cap without being per request.""" + event_loop: EventLoopStalls | None = None def team(self, name: str) -> Team: return next(team for team in self.teams if team.name == name) diff --git a/tests/e2e/perf/setup/podcpu.py b/tests/e2e/perf/setup/podcpu.py index c2402b38b71..400e5b7d38b 100644 --- a/tests/e2e/perf/setup/podcpu.py +++ b/tests/e2e/perf/setup/podcpu.py @@ -103,7 +103,7 @@ class NoSampler: return None -def _in_cluster() -> tuple[str, str, str] | None: +def in_cluster() -> tuple[str, str, str] | None: token: Final = _SERVICE_ACCOUNT / "token" ca: Final = _SERVICE_ACCOUNT / "ca.crt" namespace: Final = _SERVICE_ACCOUNT / "namespace" @@ -115,7 +115,7 @@ def _in_cluster() -> tuple[str, str, str] | None: @contextmanager def gateway_cpu_sampler() -> Generator[GatewayCpuSampler | NoSampler]: selector: Final = os.environ.get(SELECTOR_ENV, "") - cluster: Final = _in_cluster() + cluster: Final = in_cluster() sampler: Final = ( GatewayCpuSampler(selector, cluster[2], cluster[0], cluster[1]) if selector and cluster else NoSampler() ) diff --git a/tests/e2e/perf/setup/run.py b/tests/e2e/perf/setup/run.py index 9ba9b45104d..272c64754c4 100644 --- a/tests/e2e/perf/setup/run.py +++ b/tests/e2e/perf/setup/run.py @@ -9,9 +9,11 @@ import shutil import subprocess import sys import tempfile +import time +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import BinaryIO, Final from e2e_config import PROXY_BASE_URL from pydantic import TypeAdapter @@ -73,7 +75,59 @@ def _keep_records_artifact(records_path: Path) -> None: _ = shutil.copyfile(records_path, Path(artifact_dir) / "perf-records.jsonl") -def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup) -> RunOutcome: +class RecordStream: + def __init__(self, source: BinaryIO) -> None: + self.source: Final = source + self.partial = b"" + + def read(self) -> tuple[WireRecord, ...]: + chunks: Final = (self.partial + self.source.read()).split(b"\n") + self.partial = chunks[-1] + return tuple(_RECORD_ADAPTER.validate_json(line) for line in chunks[:-1] if line) + + +def run_until( + command: list[str], + env: dict[str, str], + records_path: Path, + stop_path: Path, + stop_when: Callable[[tuple[WireRecord, ...]], bool], + timeout: int, +) -> subprocess.CompletedProcess[str]: + with records_path.open("rb") as records, records_path.with_suffix(".log").open("w+") as output: + with subprocess.Popen(command, env=env, stdout=output, stderr=subprocess.STDOUT, text=True) as process: + deadline: Final = time.monotonic() + timeout + reader: Final = RecordStream(records) + try: + while process.poll() is None: + if stop_when(reader.read()): + stop_path.touch() + if time.monotonic() >= deadline: + raise subprocess.TimeoutExpired(command, timeout) + try: + _ = process.wait(timeout=1) + except subprocess.TimeoutExpired: + pass + finally: + stop_path.touch() + try: + _ = process.wait(timeout=35) + except subprocess.TimeoutExpired: + process.kill() + _ = process.wait(timeout=10) + _keep_records_artifact(records_path) + _ = output.seek(0) + return subprocess.CompletedProcess(command, process.returncode, output.read(), "") + + +def run_plan( + case: PerformanceTestCase, + plan: tuple[PlannedRequest, ...], + setup: CaseSetup, + *, + stop_when: Callable[[tuple[WireRecord, ...]], bool] | None = None, + request_timeout: float = 120.0, +) -> RunOutcome: """Replay `plan` with Locust in a child process and collect its records. The run-time ceiling is well past the planned duration so Locust's own clock never cuts a slow run short; the locustfile quits itself once the last record is written.""" @@ -82,6 +136,7 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: with tempfile.TemporaryDirectory(prefix="e2e-perf-") as directory: plan_path: Final = Path(directory) / "plan.json" records_path: Final = Path(directory) / "records.jsonl" + stop_path: Final = Path(directory) / "stop" _ = plan_path.write_bytes(_PLAN_ADAPTER.dump_json(wire)) records_path.touch() env: Final = { @@ -91,33 +146,33 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: "PERF_RECORDS_FILE": str(records_path), "PERF_RPS": str(case.traffic.rps), "PERF_USERS": str(case.traffic.concurrency), + "PERF_REQUEST_TIMEOUT": str(request_timeout), + **({"PERF_STOP_FILE": str(stop_path)} if stop_when is not None else {}), } - completed: Final = subprocess.run( - [ - sys.executable, - "-P", - "-m", - "locust", - "--headless", - "--only-summary", - "--locustfile", - str(_LOCUSTFILE), - "--host", - PROXY_BASE_URL, - "--users", - str(case.traffic.concurrency), - "--spawn-rate", - str(case.traffic.concurrency), - "--run-time", - f"{ceiling_s}s", - "--exit-code-on-error", - "0", - ], - env=env, - capture_output=True, - text=True, - timeout=ceiling_s + 120, - check=False, + command: Final = [ + sys.executable, + "-P", + "-m", + "locust", + "--headless", + "--only-summary", + "--locustfile", + str(_LOCUSTFILE), + "--host", + PROXY_BASE_URL, + "--users", + str(case.traffic.concurrency), + "--spawn-rate", + str(case.traffic.concurrency), + "--run-time", + f"{ceiling_s}s", + "--exit-code-on-error", + "0", + ] + completed: Final = ( + run_until(command, env, records_path, stop_path, stop_when, ceiling_s + 120) + if stop_when is not None + else subprocess.run(command, env=env, capture_output=True, text=True, timeout=ceiling_s + 120, check=False) ) records: Final = tuple( _RECORD_ADAPTER.validate_json(line) diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py new file mode 100644 index 00000000000..37e5c7bc904 --- /dev/null +++ b/tests/e2e/perf/setup/upgrade.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import psycopg +from cases.small_bodies_5rps import CASE +from e2e_config import MASTER_KEY +from e2e_http import send_text_external +from prometheus_client.parser import text_string_to_metric_families +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from setup.performance_test_case import PerformanceTestCase +from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes +from setup.wire import WireRecord + +MAX_TRAFFIC_SECONDS: Final = 6 * 3600 +_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...]) +_ACTIVITY: Final = TypeAdapter(tuple[list[int], int]) +_COUNT: Final = TypeAdapter(tuple[int]) +_SERVICES: Final = ("postgres", "batch_write_to_db") + + +def upgrade_case() -> PerformanceTestCase: + return replace( + CASE, + name="database_upgrade_5rps", + traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)), + limits=tuple( + limit + for limit in CASE.limits + if limit.metric + in { + "unexpected_status_count", + "gateway_added_seconds", + "time_to_first_token_seconds", + "memory_growth_megabytes", + } + ), + ) + + +class MetricsHeaders(BaseModel): + model_config = ConfigDict(frozen=True) + authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization") + + +def postgres_failures(text: str) -> float: + families: Final = tuple(text_string_to_metric_families(text)) + names: Final = frozenset(family.name for family in families) + required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES) + samples: Final = tuple(sample for family in families for sample in family.samples) + has_database_activity: Final = any( + sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0 + for sample in samples + ) + assert required <= names or has_database_activity, f"Missing database service counters: {required - names}" + failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required}) + assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter" + return sum(failures) + + +def assert_database_metrics(urls: tuple[str, ...]) -> None: + assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint" + for url in urls: + assert_database_metrics_url(url) + + +def assert_database_metrics_url(url: str) -> None: + result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10) + assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}" + assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}" + + +def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]: + ids: Final = tuple(record.call_id for record in records if record.call_id) + if not ids: + return frozenset() + rows: Final = _ROWS.validate_python( + connection.execute( + 'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),) + ).fetchall() + ) + return matched_response_ids(records, rows) + + +def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]: + call_ids: Final = frozenset(call_id for _, call_id in rows if call_id) + return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids) + + +def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]: + return _ACTIVITY.validate_python( + connection.execute("""SELECT + ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()), + (SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock') + """).fetchone() + ) + + +def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None: + assert state.phase == "complete", f"Upgrade did not finish: {state.phase}" + assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid" + (missing,) = _COUNT.validate_python( + connection.execute( + """ + SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS ( + SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name + AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""", + (list(state.required_migrations),), + ).fetchone() + ) + assert missing == 0, f"{missing} target migrations were not applied" + (unfinished,) = _COUNT.validate_python( + connection.execute(""" + SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL + """).fetchone() + ) + assert unfinished == 0, f"{unfinished} unfinished migrations" + assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, ( + "Baseline rows changed or disappeared" + ) + + +class UpgradeMonitor: + def __init__( + self, + *, + state: Callable[[], UpgradeState], + check_metrics: Callable[[], None], + read_activity: Callable[[], tuple[list[int], int]], + read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]], + start_indexes: Callable[[], None], + clock: Callable[[], float] = time.monotonic, + wall_clock: Callable[[], float] = time.time, + ) -> None: + self.state: Final = state + self.check_metrics: Final = check_metrics + self.read_activity: Final = read_activity + self.read_persisted: Final = read_persisted + self.start_indexes: Final = start_indexes + self.clock: Final = clock + self.wall_clock: Final = wall_clock + self.started: Final = clock() + self.status_limit: Final = next(limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count") + self.pending: tuple[WireRecord, ...] = () + self.processed = 0 + self.successes = 0 + self.persisted_count = 0 + self.overlap_writes = 0 + self.max_lock_waiters = 0 + self.previous_builds: frozenset[int] = frozenset() + self.build_seen_at: float | None = None + self.released = False + self.completed_at: float | None = None + + def poll(self, records: tuple[WireRecord, ...]) -> bool: + self.check_metrics() + current: Final = self.state() + assert current.phase != "failed", "Target migration/index worker failed" + assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout" + successful: Final = tuple(record for record in records if record.status == 200) + self.processed += len(records) + self.successes += len(successful) + assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses" + assert all(record.ok and record.text for record in successful), "Failed or empty inference response" + assert all(record.response_id and record.call_id for record in successful), ( + "Response missing spend-log identifiers" + ) + builds, lock_waiters = self.read_activity() + if not self.previous_builds.intersection(builds): + self.build_seen_at = self.wall_clock() if builds else None + self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters) + self.pending += successful + found: Final = self.read_persisted(self.pending) + self.persisted_count += len(found) + still_building, _ = self.read_activity() + since: Final = self.build_seen_at + if since is not None and frozenset(builds).intersection(still_building): + self.overlap_writes += sum( + 1 for record in self.pending if record.response_id in found and record.started_at >= since + ) + self.previous_builds = frozenset(builds) + self.pending = tuple(record for record in self.pending if record.response_id not in found) + if self.successes and not self.released: + self.start_indexes() + self.released = True + if current.phase == "complete" and self.completed_at is None: + self.completed_at = self.clock() + return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300 + + def assert_complete(self, state: UpgradeState) -> None: + assert self.successes >= 300, f"Only {self.successes} requests completed" + assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log" + assert self.persisted_count == self.successes, "Spend-log count does not match successful requests" + assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work" + + +def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor: + return UpgradeMonitor( + state=lambda: read_state(connection), + check_metrics=lambda: assert_database_metrics(urls), + read_activity=lambda: activity(connection), + read_persisted=lambda records: persisted(connection, records), + start_indexes=lambda: release_indexes(connection), + ) diff --git a/tests/e2e/perf/setup/upgrade_worker.py b/tests/e2e/perf/setup/upgrade_worker.py new file mode 100644 index 00000000000..e69e841814f --- /dev/null +++ b/tests/e2e/perf/setup/upgrade_worker.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import json +import os +import sys +import time +from pathlib import Path +from typing import Final, Literal + +import psycopg +from pydantic import BaseModel, ConfigDict, TypeAdapter + + +class UpgradeState(BaseModel): + model_config = ConfigDict(frozen=True) + + phase: Literal["schema", "schema_ready", "indexes", "complete", "failed"] + started_at: float + schema_seconds: float = 0 + index_seconds: float = 0 + required_indexes: tuple[str, ...] + required_migrations: tuple[str, ...] + pending_indexes: tuple[str, ...] + baseline_rows: dict[str, str] + + +_STATE: Final = TypeAdapter(tuple[UpgradeState]) +_BOOL: Final = TypeAdapter(tuple[bool]) +_STRINGS: Final = TypeAdapter(tuple[tuple[str], ...]) +_FINGERPRINTS: Final = TypeAdapter(tuple[tuple[str, str], ...]) +_FINGERPRINT: Final = """md5(ROW(request_id, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, model, "startTime", "endTime")::text)""" + + +def connect() -> psycopg.Connection[tuple[object, ...]]: + return psycopg.connect( + os.environ["DATABASE_URL"], + autocommit=True, + connect_timeout=10, + options="-c statement_timeout=10000 -c timezone=UTC", + ) + + +def fingerprints(connection: psycopg.Connection[tuple[object, ...]], ids: tuple[str, ...]) -> dict[str, str]: + rows: Final = connection.execute( + f'SELECT request_id, {_FINGERPRINT} FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', (list(ids),) + ).fetchall() + return dict(_FINGERPRINTS.validate_python(rows)) + + +def missing_indexes(connection: psycopg.Connection[tuple[object, ...]], names: tuple[str, ...]) -> tuple[str, ...]: + rows: Final = connection.execute( + """SELECT name FROM unnest(%s::text[]) AS name WHERE NOT EXISTS ( + SELECT 1 FROM pg_index WHERE indexrelid = to_regclass(format('public.%%I', name)) + AND indisvalid AND indisready)""", + (list(names),), + ).fetchall() + return tuple(row[0] for row in _STRINGS.validate_python(rows)) + + +def read_state(connection: psycopg.Connection[tuple[object, ...]]) -> UpgradeState: + (state,) = _STATE.validate_python( + connection.execute('SELECT state FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone() + ) + return state + + +def write_state(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None: + _ = connection.execute( + """INSERT INTO "_LiteLLMPerfUpgrade" (id, state) VALUES (1, %s::jsonb) + ON CONFLICT (id) DO UPDATE SET state = EXCLUDED.state""", + (state.model_dump_json(),), + ) + + +def release_indexes(connection: psycopg.Connection[tuple[object, ...]]) -> None: + _ = connection.execute('UPDATE "_LiteLLMPerfUpgrade" SET traffic_ready=true WHERE id=1') + + +def wait_for_traffic(connection: psycopg.Connection[tuple[object, ...]], deadline: float) -> bool: + while time.monotonic() < deadline: + (ready,) = _BOOL.validate_python( + connection.execute('SELECT traffic_ready FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone() + ) + if ready: + return True + time.sleep(1) + return False + + +def main() -> int: + assert os.environ.get("E2E_DATABASE_UPGRADE") == "1", "Only run against a disposable restored baseline" + assert os.environ.get("DATABASE_PERF_CONFIRM_DISPOSABLE") == "1", "Disposable database confirmation required" + from litellm_proxy_extras import request_log_indexes + from litellm_proxy_extras.utils import ProxyExtrasDBManager + + with connect() as connection: + _ = connection.execute("""CREATE TABLE "_LiteLLMPerfUpgrade" ( + id integer PRIMARY KEY CHECK (id=1), state jsonb NOT NULL, traffic_ready boolean NOT NULL DEFAULT false)""") + sample: Final = connection.execute(""" + (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id LIMIT 32) + UNION (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id DESC LIMIT 32)""").fetchall() + ids: Final = tuple(row[0] for row in _STRINGS.validate_python(sample)) + assert ids, "Restored baseline has no spend logs" + names: Final = tuple(index.name for index in request_log_indexes.REQUEST_LOG_INDEXES) + migrations: Final = Path(request_log_indexes.__file__).parent / "migrations" + initial: Final = UpgradeState( + phase="schema", + started_at=time.time(), + required_indexes=names, + required_migrations=tuple(sorted(path.parent.name for path in migrations.glob("*/migration.sql"))), + pending_indexes=missing_indexes(connection, names), + baseline_rows=fingerprints(connection, ids), + ) + assert initial.required_migrations and initial.required_indexes, "Target schema requirements are missing" + write_state(connection, initial) + try: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True), ( + "Schema migration failed" + ) + ready: Final = initial.model_copy( + update={ + "phase": "schema_ready", + "schema_seconds": time.time() - initial.started_at, + "pending_indexes": missing_indexes(connection, names), + } + ) + write_state(connection, ready) + assert wait_for_traffic(connection, time.monotonic() + 1800), "Traffic did not start within 30 minutes" + write_state(connection, ready.model_copy(update={"phase": "indexes"})) + index_start: Final = time.monotonic() + assert ProxyExtrasDBManager.build_request_log_indexes(), "Request-log index build failed" + complete: Final = ready.model_copy( + update={"phase": "complete", "index_seconds": time.monotonic() - index_start} + ) + write_state(connection, complete) + _ = sys.stdout.write(complete.model_dump_json() + "\n") + return 0 + except Exception as exc: + write_state(connection, read_state(connection).model_copy(update={"phase": "failed"})) + _ = sys.stderr.write(json.dumps({"error": type(exc).__name__}) + "\n") + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/e2e/perf/setup/wire.py b/tests/e2e/perf/setup/wire.py index 07dc5d170ae..ced79d6970c 100644 --- a/tests/e2e/perf/setup/wire.py +++ b/tests/e2e/perf/setup/wire.py @@ -9,12 +9,13 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final, Literal -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from setup.corpus.session import Endpoint as CorpusEndpoint from setup.corpus.session import request_body as corpus_body Endpoint = Literal["chat_completions", "responses", "messages"] +_PAYLOAD: Final = TypeAdapter(dict[str, object]) class TokenSource(BaseModel): @@ -70,6 +71,33 @@ class WireRecord(BaseModel): mock_delay_s: float | None call_id: str | None error: str | None + response_id: str | None = None + + +def response_id(payload: Mapping[str, object]) -> str | None: + nested: Final = payload.get("response") or payload.get("message") + source: Final = _PAYLOAD.validate_python(nested) if isinstance(nested, Mapping) else payload + identifier: Final = source.get("id") + return identifier if isinstance(identifier, str) else None + + +class PlanProgress: + def __init__(self, total: int) -> None: + self.total: Final = total + self.issued = 0 + self.completed = 0 + self.stopped = False + + def take(self, *, stop: bool) -> bool: + self.stopped = self.stopped or stop + if self.stopped or self.issued >= self.total: + return False + self.issued += 1 + return True + + @property + def drained(self) -> bool: + return (self.stopped or self.issued == self.total) and self.completed == self.issued _CORPUS_ENDPOINT: Final[Mapping[Endpoint, CorpusEndpoint]] = MappingProxyType( diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py new file mode 100644 index 00000000000..f165b9b738f --- /dev/null +++ b/tests/e2e/perf/test_database_upgrade.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path +from typing import Final + +import pytest +from idp import Keycloak +from lifecycle import ResourceManager +from memory_readings import read_rss_everywhere +from proxy_client import ProxyClient +from setup.datastore import read_db_counters, read_redis_counters +from setup.observe import Observation, evaluate, read_mock_requests, render, reset_mock, set_mock_latency +from setup.plan import build_plan +from setup.podcpu import gateway_cpu_sampler +from setup.provision import MOCK_PROVIDERS, build_setup, mock_url +from setup.run import run_plan +from setup.upgrade import assert_finished, monitor_for, upgrade_case +from setup.upgrade_worker import connect, read_state + +pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] + + +@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only") +def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None: + case: Final = upgrade_case() + urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS) + for url in urls: + _ = reset_mock(url) + _ = set_mock_latency(url, case.mock_latency) + setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources) + metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip()) + with connect() as connection: + assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting" + rss_before: Final = read_rss_everywhere(proxy) + db_before: Final = read_db_counters() + redis_before: Final = read_redis_counters() + monitor: Final = monitor_for(connection, metrics) + with gateway_cpu_sampler() as cpu: + outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30) + assert outcome.locust_exit == 0, outcome.locust_output + remaining: Final = outcome.records[monitor.processed :] + _ = monitor.poll(remaining) + deadline: Final = time.monotonic() + 120 + while monitor.pending and time.monotonic() < deadline: + time.sleep(1) + _ = monitor.poll(()) + state: Final = read_state(connection) + _ = sys.stdout.write( + f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} " + f"unmatched_requests={len(monitor.pending)}\n" + ) + monitor.assert_complete(state) + assert_finished(connection, state) + rss_after: Final = read_rss_everywhere(proxy) + assert {r.worker for r in rss_before.readings} == {r.worker for r in rss_after.readings}, ( + "Gateway workers restarted" + ) + observed: Final = Observation( + records=outcome.records, + mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)), + guardrail_calls=(), + rss_before=rss_before, + rss_after=rss_after, + wall_seconds=outcome.wall_seconds, + spend_logged_share=1, + db_before=db_before, + db_after=read_db_counters(), + redis_before=redis_before, + redis_after=read_redis_counters(), + gateway_cpu_seconds=cpu.seconds, + ) + rows, exceeded = evaluate(case, observed) + _ = sys.stdout.write(render(case, observed, rows) + "\n") + _ = sys.stdout.write( + f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} " + f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} " + f"pending_indexes={state.pending_indexes}\n" + ) + artifact: Final = os.environ.get("E2E_ARTIFACT_DIR") + if artifact: + _ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8") + assert not exceeded, "\n".join(exceeded) diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index d89a8634ab3..902483f8a2f 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -5,7 +5,9 @@ clients saw, what the mocks received and how the workers' memory moved.""" from __future__ import annotations +import os import time +from datetime import UTC, datetime from typing import Final import pytest @@ -16,6 +18,16 @@ from memory_readings import read_rss_everywhere from models import SpendLogsParams from proxy_client import ProxyClient from setup.datastore import read_db_counters, read_redis_counters, redis_monitor +from setup.event_loop import ( + LOG_FILE_ENV, + OPT_IN_ENV, + StallRow, + offenders, + parse_stalls, + read_gateway_log, + render_stalls, + summarize, +) from setup.observe import ( Observation, evaluate, @@ -48,6 +60,17 @@ def _spend_logged_share(proxy: ProxyClient, call_ids: tuple[str, ...]) -> float return found / len(call_ids) +def _event_loop_rows(case: PerformanceTestCase, since: datetime) -> tuple[StallRow, ...]: + """The gateway's asyncio debug holds for the traffic window; a case that asks for them cannot pass without a log.""" + assert case.event_loop is not None + lines: Final = read_gateway_log(since) + assert lines is not None, ( + f"{case.name} needs the gateway log: run in the perf lane (pods/log with the runner service account) or set " + f"{LOG_FILE_ENV} to the gateway's log file" + ) + return summarize(case.event_loop, parse_stalls(lines)) + + @pytest.mark.parametrize("case", CASES, ids=[case.name for case in CASES]) def test_perf_case( case: PerformanceTestCase, @@ -56,6 +79,11 @@ def test_perf_case( resources: ResourceManager, request: pytest.FixtureRequest, ) -> None: + if case.event_loop is not None and not os.environ.get(OPT_IN_ENV): + pytest.skip( + f"{case.name} reads asyncio debug output: deploy the gateway with PYTHONASYNCIODEBUG=1 and " + f"LITELLM_EVENT_LOOP=asyncio, then set {OPT_IN_ENV}=1" + ) plan: Final = build_plan(case) mock_urls: Final = {provider: mock_url(provider) for provider in MOCK_PROVIDERS} for url in (*mock_urls.values(), guardrail_url()): @@ -68,6 +96,7 @@ def test_perf_case( db_before: Final = read_db_counters() redis_before: Final = read_redis_counters() run_started: Final = time.monotonic() + traffic_started_at: Final = datetime.now(tz=UTC) with gateway_cpu_sampler() as cpu, redis_monitor() as monitor: outcome: Final = run_plan(case, plan, setup) run_finished: Final = time.monotonic() @@ -102,6 +131,7 @@ def test_perf_case( gateway_cpu_seconds=cpu.seconds, ) rows, exceeded_limits = evaluate(case, observation) + stall_rows: Final = _event_loop_rows(case, traffic_started_at) if case.event_loop is not None else () item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped assert isinstance(item, pytest.Item) item.user_properties[:] = [ @@ -109,6 +139,8 @@ def test_perf_case( ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set item.user_properties.extend(junit_properties(rows)) print(render(case, observation, rows)) + if case.event_loop is not None: + print(render_stalls(case.event_loop, stall_rows)) print( f" phases: setup {run_started - setup_started:.0f}s, locust {run_finished - run_started:.0f}s, " f"observe {time.monotonic() - run_finished:.0f}s (rss {rss_read - run_finished:.0f}s, spend sample " @@ -118,3 +150,4 @@ def test_perf_case( f"{len(plan)} planned, {len(outcome.records)} recorded; locust said:\n{outcome.locust_output}" ) assert not exceeded_limits, "limits exceeded (observed > max):\n" + "\n".join(exceeded_limits) + assert not offenders(stall_rows), "event-loop holds over the threshold:\n" + "\n".join(offenders(stall_rows)) diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/router/test_reliability_memory_e2e.py index 9ae6edb3eb2..196784b5345 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/router/test_reliability_memory_e2e.py @@ -13,19 +13,26 @@ deployment that fails a handful of times in a row is benched (a bad-credential 4 included), the router answers "No deployments available" without retrying, and the retry loop that leaks stops running; cooldown_time 0 keeps it running. -Two identical phases run back to back. The first is the warmup that grows the -proxy's caches and allocator arenas to their steady state, the second is the one -the budget applies to, so a healthy proxy shows the second phase adding roughly -nothing while a leaking one adds a fixed amount per request. RSS is read through -/debug/memory/summary on every configured replica; a burst of failing calls leaves -a transient bulge of garbage that gc reclaims within seconds, so each checkpoint -samples until no new worker has answered for a settle window and keeps the lowest -reading per worker. The growth is judged per worker (by replica address, hostname -and pid, since pods in their own pid namespaces report the same pids) so each -worker is compared with itself, and the two checkpoints must see the same workers: -a single load-balanced address reaches the workers behind it one answer at a time, -and a worker that answered only one checkpoint would otherwise drop out of the -comparison, which is where a leaking worker could hide. +Three identical batches run back to back. The first two are the warmup that grows +the proxy's caches and allocator arenas to their steady state, the third is the one +the budget applies to, so a healthy proxy shows the third batch adding roughly +nothing while a leaking one adds a fixed amount per request. One warmup batch is not +enough: each batch's in-flight retries peak well above the live heap, the allocator +keeps that high-water mark, and on a cold worker the second batch still set a new one +(about 50 MB of RSS with the live Python heap flat to within 1 MB), where the later +batches each settled under 25 MB. RSS is read through /debug/memory/summary on every +configured replica; a burst of failing calls leaves a transient bulge of garbage that +gc reclaims within seconds, so each checkpoint samples until no new worker has +answered for a settle window and keeps the lowest reading per worker. A checkpoint +runs before the warmup and after every batch, and each batch's settled growth per +worker rides out as a repeated `rss_batch_growth` JUnit property as soon as that +batch settles, so a run that fails keeps the curve up to the batch it failed in. The budget is judged per worker +(by replica address, hostname and pid, since pods in their own pid namespaces report +the same pids) so each worker is compared with itself, and the checkpoints around the +measured batch must see the same workers: a single load-balanced address reaches the +workers behind it one answer at a time, and a worker that answered only one +checkpoint would otherwise drop out of the comparison, which is where a leaking +worker could hide. RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS @@ -58,7 +65,7 @@ from __future__ import annotations import json import time -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from types import MappingProxyType @@ -91,6 +98,8 @@ pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] DEPLOYMENTS_PER_GROUP: Final = 2 RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES +WARMUP_BATCHES: Final = 2 +MEASURED_BATCH: Final = WARMUP_BATCHES + 1 @dataclass(frozen=True, slots=True) @@ -177,12 +186,52 @@ def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[WorkerKey, RssReading ) +def _record_batch_growth( + record_property: Callable[[str, object], None], + batch: int, + before: Mapping[WorkerKey, RssReading], + after: Mapping[WorkerKey, RssReading], +) -> None: + for growth in (WorkerGrowth(before[worker], after[worker]) for worker in before.keys() & after.keys()): + record_property( + "rss_batch_growth", + json.dumps( + { + "batch": batch, + "worker": growth.warm.where, + "before_mb": growth.warm.ram_usage_mb, + "after_mb": growth.after.ram_usage_mb, + "growth_mb": round(growth.growth_mb, 2), + } + ), + ) + + +def _settled_rss_around_measured_batch( + proxy: ProxyClient, + key: str, + model: str, + override: RouterSettingsOverride, + fallback: str, + record_property: Callable[[str, object], None], + before: Mapping[WorkerKey, RssReading], + batch: int, +) -> tuple[Mapping[WorkerKey, RssReading], Mapping[WorkerKey, RssReading]]: + calls: Final = _fail_many(proxy, key, model, override) + _assert_every_call_failed_through_fallback(calls, fallback) + after: Final = _settled_rss_per_worker(proxy) + _record_batch_growth(record_property, batch, before, after) + if batch == MEASURED_BATCH: + return before, after + return _settled_rss_around_measured_batch(proxy, key, model, override, fallback, record_property, after, batch + 1) + + def _heaviest_worker_growth( warm: Mapping[WorkerKey, RssReading], after: Mapping[WorkerKey, RssReading] ) -> WorkerGrowth: assert warm.keys() == after.keys(), ( - f"the workers answering /debug/memory/summary changed between the checkpoints, so not every worker can " - f"be compared with itself: gone after the measured batch {sorted(warm.keys() - after.keys())} (a worker " + f"the workers answering /debug/memory/summary changed around measured batch {MEASURED_BATCH}, so not every " + f"worker can be compared with itself: gone after it {sorted(warm.keys() - after.keys())} (a worker " f"that died or was restarted under failing traffic, which is what an OOM kill looks like), first seen " f"after it {sorted(after.keys() - warm.keys())} (the warm window never reached them, so they have no " f"baseline; raise E2E_MEMORY_RSS_SETTLE_SAMPLES if the stack has more workers than the window covers)" @@ -247,6 +296,7 @@ class TestReliabilityMemory: f"leaves the worker little room for traffic" ) + @pytest.mark.filterwarnings("ignore:record_property is incompatible with junit_family") @pytest.mark.covers("reliability.perf.memory.under_slo") @meta( Subject( @@ -257,7 +307,11 @@ class TestReliabilityMemory: ) ) def test_failing_requests_do_not_grow_rss_or_stored_request( - self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + self, + client: ComplexityRouterClient, + resources: ResourceManager, + scoped_key: str, + record_property: Callable[[str, object], None], ) -> None: marker: Final = unique_marker() primary: Final = f"reliability-memory-{marker}" @@ -278,19 +332,23 @@ class TestReliabilityMemory: f"request into the stored snapshot the way the v1.100.0 ones did" ) - warmup: Final = _fail_many(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback(warmup, fallback) - warm: Final = _settled_rss_per_worker(client.proxy) - - measured: Final = _fail_many(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback(measured, fallback) - after: Final = _settled_rss_per_worker(client.proxy) + warm, after = _settled_rss_around_measured_batch( + client.proxy, + scoped_key, + primary, + override, + fallback, + record_property, + _settled_rss_per_worker(client.proxy), + 1, + ) heaviest: Final = _heaviest_worker_growth(warm, after) assert heaviest.growth_mb <= MEMORY_RSS_BUDGET_MB, ( - f"proxy RSS grew {heaviest.growth_mb:.1f} MB over a second batch of {MEMORY_REQUESTS_PER_PHASE} failing " - f"requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) after an identical warmup batch, " - f"past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: worker pid {heaviest.warm.worker_pid} on " + f"proxy RSS grew {heaviest.growth_mb:.1f} MB over measured batch {MEASURED_BATCH} of " + f"{MEMORY_REQUESTS_PER_PHASE} failing requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) " + f"after {WARMUP_BATCHES} identical warmup batches, past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: " + f"worker pid {heaviest.warm.worker_pid} on " f"{heaviest.warm.hostname or 'an unnamed host'} behind {heaviest.warm.replica} settled at " f"{heaviest.warm.ram_usage_mb:.1f} MB warm and " f"{heaviest.after.ram_usage_mb:.1f} MB after; failing requests are leaking memory the way the " diff --git a/tests/e2e/router/test_reliability_prompt_caching_e2e.py b/tests/e2e/router/test_reliability_prompt_caching_e2e.py index 7bf23bf967f..25e842f362e 100644 --- a/tests/e2e/router/test_reliability_prompt_caching_e2e.py +++ b/tests/e2e/router/test_reliability_prompt_caching_e2e.py @@ -43,6 +43,7 @@ FOLLOW_UPS = 3 class TestReliabilityPromptCachingAffinity: + @pytest.mark.provider_live @pytest.mark.covers("reliability.cache.prompt_caching_model_select.returns_cached") @meta( Subject( diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts index 50850a5cd58..d0b1bb9ecfb 100644 --- a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -18,8 +18,11 @@ const test = base.extend<{ scopedKey: ScopedKey }>({ team_id: E2E_TEAM_CRUD_ID, models: [CHAT_MODEL_A], }); - await use({ alias, token: created.token }); - await deleteVirtualKey(page.request, created.token); + try { + await use({ alias, token: created.token }); + } finally { + await deleteVirtualKey(page.request, created.token); + } }, }); @@ -36,15 +39,44 @@ test.describe("Proxy Admin - Key budget window", () => { await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); - await openKeyDetail(page, alias); - - await page.getByRole("tab", { name: "Settings" }).click(); - await page.getByRole("button", { name: "Edit Settings" }).click(); - - await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + // The initial, unfiltered list must already contain this newly created key. + // useKeys keeps that list visible while the filtered search is pending. + await expect(page.getByRole("row").filter({ hasText: alias })).toBeVisible({ timeout: 15_000 }); + // Deliver the real search result after editing starts: replacing the list's + // key object must not reset the open form and erase its unsaved budget. + const searchReceived = Promise.withResolvers(); + const releaseSearch = Promise.withResolvers(); + const isAliasSearch = (url: string) => { + const parsed = new URL(url); + return parsed.pathname.endsWith("/key/list") && parsed.searchParams.get("search") === alias; + }; + await page.route("**/key/list?**", async (route) => { + if (!isAliasSearch(route.request().url())) return route.continue(); + const response = await route.fetch(); + searchReceived.resolve(); + await releaseSearch.promise; + await route.fulfill({ response }); + }); + try { + await openKeyDetail(page, alias); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await searchReceived.promise; + const refreshed = page.waitForResponse((response) => isAliasSearch(response.url())); + releaseSearch.resolve(); + await refreshed; + } finally { + releaseSearch.resolve(); + await page.unrouteAll({ behavior: "wait" }); + } await page.getByLabel("Reset Budget", { exact: true }).click(); await page.getByRole("option", { name: "monthly", exact: true }).click(); - await page.getByRole("button", { name: "Save Changes" }).click(); + const saved = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(saved.max_budget, "a delayed search response must not erase the budget draft").toBe("12.5"); + expect(saved.budget_duration).toBe("30d"); await expect .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { diff --git a/tests/e2e_harness/perf/setup/test_event_loop.py b/tests/e2e_harness/perf/setup/test_event_loop.py new file mode 100644 index 00000000000..57868305716 --- /dev/null +++ b/tests/e2e_harness/perf/setup/test_event_loop.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from typing import Final + +from setup.event_loop import EventLoopStalls, KnownStall, offenders, parse_stalls, summarize + +_HANDLER: Final = ( + "Executing wait_for= created at /usr/local/lib/python3.12/asyncio/base_events.py:448> took 0.731 seconds" +) +_STREAM: Final = ( + "Executing created at " + "/usr/local/lib/python3.12/asyncio/tasks.py:420> took 0.412 seconds" +) +_STARTUP: Final = ( + "Executing created at " + "/usr/local/lib/python3.12/asyncio/base_events.py:448> took 0.302 seconds" +) +_GC_PASS: Final = ( + "2026-10-17 04:46:02,101 - LiteLLM Proxy - WARNING - Garbage collection gen=2 collected=98608 took 0.711 seconds" +) +_SPEC: Final = EventLoopStalls( + allowlist=( + KnownStall("LifespanOn.main", max_count=3), + KnownStall("RequestResponseCycle.run_asgi", max_count=2), + KnownStall("garbage collection", max_count=1), + ) +) + + +def test_parse_keeps_the_coroutine_its_location_and_the_hold() -> None: + (stall,) = parse_stalls(("INFO: plain log line", _HANDLER)) + assert stall.coroutine == "RequestResponseCycle.run_asgi()" + assert stall.where == "uvicorn/protocols/http/h11_impl.py:416" + assert stall.seconds == 0.731 + + +def test_a_location_off_the_allowlist_offends_once_however_often_it_appears() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, _STREAM, _HANDLER))) + assert [(row.coroutine, row.count, row.offending) for row in rows] == [ + ("TaskHandle._run_coro()", 2, True), + ("RequestResponseCycle.run_asgi()", 1, False), + ("LifespanOn.main()", 1, False), + ] + assert offenders(rows) == ( + "TaskHandle._run_coro() anyio/_core/_tasks.py:275: 2x, max 0.41 s (not on the allowlist)", + ) + + +def test_a_known_location_offends_only_past_its_count() -> None: + within: Final = summarize(_SPEC, parse_stalls((_HANDLER, _HANDLER))) + assert offenders(within) == () + past: Final = summarize(_SPEC, parse_stalls((_HANDLER, _HANDLER, _HANDLER))) + assert offenders(past) == ( + "RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416: 3x, max 0.73 s " + "(allowlist RequestResponseCycle.run_asgi permits 2)", + ) + + +def test_holds_under_the_threshold_are_not_rows() -> None: + short: Final = _HANDLER.replace("took 0.731 seconds", "took 0.099 seconds") + assert summarize(_SPEC, parse_stalls((short,))) == () + + +def test_a_hold_the_collector_pass_before_it_accounts_for_is_the_collectors_not_the_tasks() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_GC_PASS, _STREAM.replace("0.412", "0.731")))) + assert [(row.coroutine, row.where, row.max_seconds, row.offending) for row in rows] == [ + ("garbage collection", "gen=2", 0.731, False) + ] + + +def test_a_hold_well_past_its_collector_pass_stays_on_the_task() -> None: + short_pass: Final = _GC_PASS.replace("0.711", "0.200") + rows: Final = summarize(_SPEC, parse_stalls((short_pass, _STREAM.replace("0.412", "0.731")))) + assert [(row.coroutine, row.offending) for row in rows] == [("TaskHandle._run_coro()", True)] + + +def test_a_hold_is_charged_to_the_longest_collector_pass_that_fits_it_not_to_their_sum() -> None: + thread_side_pass: Final = _GC_PASS.replace("0.711", "0.627") + rows: Final = summarize( + _SPEC, parse_stalls((thread_side_pass, _GC_PASS.replace("0.711", "0.605"), _STREAM.replace("0.412", "0.614"))) + ) + assert [(row.coroutine, row.where, row.max_seconds, row.offending) for row in rows] == [ + ("garbage collection", "gen=2", 0.614, False) + ] + + +def test_a_collector_pass_explains_only_the_next_hold_and_only_one_at_least_as_long() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_GC_PASS, _STARTUP, _HANDLER))) + assert sorted((row.coroutine, row.offending) for row in rows) == [ + ("LifespanOn.main()", False), + ("RequestResponseCycle.run_asgi()", False), + ] + assert summarize(_SPEC, parse_stalls((_GC_PASS,))) == () diff --git a/tests/e2e_harness/perf/test_run.py b/tests/e2e_harness/perf/test_run.py new file mode 100644 index 00000000000..23f62c9c560 --- /dev/null +++ b/tests/e2e_harness/perf/test_run.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Final + +import pytest +from setup.run import run_until +from setup.wire import WireRecord +from test_upgrade import record + + +class StopAfterRecord: + def __init__(self, *, partial: Path, resume: Path, fail: bool = False) -> None: + self.partial: Final = partial + self.resume: Final = resume + self.fail: Final = fail + self.received: tuple[WireRecord, ...] = () + + def __call__(self, records: tuple[WireRecord, ...]) -> bool: + self.received += records + if self.partial.exists(): + self.resume.touch() + assert not (self.fail and records), "Database failed" + return bool(records) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_controlled_process_preserves_partial_lines_and_drains_on_stop_or_failure(tmp_path: Path, fail: bool) -> None: + records: Final = tmp_path / "records.jsonl" + stop: Final = tmp_path / "stop" + partial: Final = tmp_path / "partial" + resume: Final = tmp_path / "resume" + records.touch() + encoded: Final = record().model_dump_json() + "\n" + program: Final = f""" +import time +from pathlib import Path +with Path({str(records)!r}).open("w") as f: + text = {encoded!r} + f.write(text[:10]) + f.flush() + Path({str(partial)!r}).touch() + while not Path({str(resume)!r}).exists(): + time.sleep(0.01) + f.write(text[10:]) + f.flush() + while not Path({str(stop)!r}).exists(): + time.sleep(0.01) + f.write({record(1).model_dump_json()!r} + "\\n") + f.flush() +""" + callback: Final = StopAfterRecord(partial=partial, resume=resume, fail=fail) + command: Final = [sys.executable, "-c", program] + if fail: + with pytest.raises(AssertionError, match="Database failed"): + _ = run_until(command, dict(os.environ), records, stop, callback, timeout=10) + else: + result: Final = run_until(command, dict(os.environ), records, stop, callback, timeout=10) + assert result.returncode == 0, result.stdout + assert callback.received == (record(),) + assert tuple(WireRecord.model_validate_json(line) for line in records.read_text().splitlines()) == ( + record(), + record(1), + ) diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py new file mode 100644 index 00000000000..0904b7a7b5a --- /dev/null +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -0,0 +1,325 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +import pytest +from cases.small_bodies_5rps import CASE +from memory_readings import RssCapture, RssReading +from setup.observe import Observation, evaluate +from setup.upgrade import MAX_TRAFFIC_SECONDS, UpgradeMonitor, matched_response_ids, postgres_failures, upgrade_case +from setup.upgrade_worker import UpgradeState +from setup.wire import PlanProgress, WireRecord, response_id + + +def record(seq: int = 0, *, started_at: float = 100) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"request-{seq}", + caller="key", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=started_at, + status=200, + header_s=1, + first_byte_s=1, + first_token_s=None, + total_s=1, + text="answer", + ok=True, + mock_delay_s=1, + call_id=f"call-{seq}", + error=None, + response_id=f"response-{seq}", + ) + + +def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None: + records: Final = (record(0), record(1), record(2)) + rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call")) + assert matched_response_ids(records, rows) == frozenset({"response-0"}) + + +def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState: + return UpgradeState.model_validate( + { + "phase": phase, + "started_at": 0, + "required_indexes": ["index"], + "required_migrations": ["migration"], + "pending_indexes": ["index"] if pending else [], + "baseline_rows": {"seed-row": "hash"}, + } + ) + + +class Probe: + def __init__(self) -> None: + self.now = 0.0 + self.current = state() + self.builds: list[int] = [] + self.found: frozenset[str] = frozenset() + self.releases = 0 + self.metric_checks = 0 + self.errors = False + + def check_metrics(self) -> None: + self.metric_checks += 1 + assert not self.errors, "Failed PostgreSQL operation" + + def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]: + return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id) + + def release(self) -> None: + self.releases += 1 + + def monitor(self) -> UpgradeMonitor: + return UpgradeMonitor( + state=lambda: self.current, + check_metrics=self.check_metrics, + read_activity=lambda: (self.builds, 2), + read_persisted=self.persisted, + start_indexes=self.release, + clock=lambda: self.now, + wall_clock=lambda: self.now, + ) + + +def test_stop_drains_issued_requests_without_issuing_the_rest() -> None: + progress: Final = PlanProgress(100) + assert progress.take(stop=False) + assert progress.take(stop=False) + progress.completed = 1 + assert not progress.take(stop=True) + assert not progress.drained + assert not progress.take(stop=False) + progress.completed = 2 + assert progress.drained + assert progress.issued == 2 + + +def test_finite_plan_still_waits_for_its_last_response() -> None: + progress: Final = PlanProgress(1) + assert progress.take(stop=False) + assert not progress.take(stop=False) + assert not progress.drained + progress.completed = 1 + assert progress.drained + + +@pytest.mark.parametrize( + "payload,expected", + [ + ({"id": "chat-id"}, "chat-id"), + ({"type": "response.created", "response": {"id": "response-id"}}, "response-id"), + ({"type": "message_start", "message": {"id": "message-id"}}, "message-id"), + ({"type": "response.output_item.added", "item": {"id": "item-id"}}, None), + ({"type": "content_block_delta", "delta": {"text": "hello"}}, None), + ], +) +def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None: + assert response_id(payload) == expected + + +def metrics(*, failure: str = "0", batch: str = "0") -> str: + return f"""# TYPE litellm_postgres_failed_requests_total counter +litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure} +# TYPE litellm_batch_write_to_db_failed_requests_total counter +litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch} +""" + + +def test_failed_sql_and_async_batch_writes_are_both_counted() -> None: + assert postgres_failures(metrics(failure="2", batch="3")) == 5 + assert postgres_failures(metrics()) == 0 + + +def test_no_error_label_series_is_a_valid_zero() -> None: + assert ( + postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter +# TYPE litellm_batch_write_to_db_failed_requests_total counter +""") + == 0 + ) + + +@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"]) +def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None: + total: Final = """# TYPE litellm_postgres_total_requests_total counter +litellm_postgres_total_requests_total{postgres="postgres"} 1 +""" + failure: Final = ( + f"# TYPE litellm_{service}_failed_requests_total counter\n" + f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n' + if service is not None + else "" + ) + assert postgres_failures(total + failure) == (1 if service is not None else 0) + + +@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"]) +def test_missing_failure_families_require_observed_database_activity(total: str) -> None: + with pytest.raises(AssertionError): + _ = postgres_failures( + "# TYPE litellm_postgres_total_requests_total counter\n" + f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n' + ) + + +@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")]) +def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None: + with pytest.raises(AssertionError): + _ = postgres_failures(body) + + +def test_index_completion_stops_after_the_post_upgrade_window() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + assert not monitor.poll((record(),)) + assert probe.releases == 1 + probe.builds = [17] + probe.now = 101 + assert not monitor.poll(()) + probe.found = frozenset({"response-0", "response-1"}) + probe.now = 103 + assert not monitor.poll((record(1, started_at=102),)) + assert monitor.overlap_writes == 1 + assert monitor.persisted_count == 2 + assert not monitor.pending + assert monitor.max_lock_waiters == 2 + probe.current = state("complete") + assert not monitor.poll(()) + probe.now = 162 + assert not monitor.poll(()) + probe.now = 163 + assert not monitor.poll(()) + probe.found = frozenset(f"response-{i}" for i in range(300)) + assert monitor.poll(tuple(record(i) for i in range(2, 300))) + assert probe.releases == 1 + assert probe.metric_checks == 7 + + +def test_http_200_without_spend_row_cannot_complete() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + records: Final = tuple(record(i) for i in range(300)) + _ = monitor.poll(records) + with pytest.raises(AssertionError, match="no matching spend log"): + monitor.assert_complete(state("complete", pending=False)) + + +def test_no_pending_index_work_still_requires_300_persisted_requests() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + records: Final = tuple(record(i) for i in range(300)) + probe.found = frozenset(f"response-{i}" for i in range(300)) + _ = monitor.poll(records) + monitor.assert_complete(state("complete", pending=False)) + assert monitor.successes == monitor.persisted_count == 300 + with pytest.raises(AssertionError, match="during index work"): + monitor.assert_complete(state("complete", pending=True)) + + +def test_postgres_failure_fails_even_when_inference_succeeded() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + probe.errors = True + with pytest.raises(AssertionError, match="PostgreSQL"): + _ = monitor.poll((record(),)) + + +def test_failed_migration_and_operational_timeout_are_not_success() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + probe.current = state("failed") + with pytest.raises(AssertionError, match="worker failed"): + _ = monitor.poll(()) + probe.current = state() + probe.now = MAX_TRAFFIC_SECONDS + with pytest.raises(AssertionError, match="operational timeout"): + _ = monitor.poll(()) + + +def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None: + case: Final = upgrade_case() + assert case.traffic.rps == 5 + assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection( + limit.metric for limit in case.limits + ) + assert case.limits == tuple( + limit + for limit in CASE.limits + if limit.metric + in { + "unexpected_status_count", + "gateway_added_seconds", + "time_to_first_token_seconds", + "memory_growth_megabytes", + } + ) + + +def failed_record(seq: int) -> WireRecord: + return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None}) + + +def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + assert not monitor.poll(tuple(failed_record(i) for i in range(3))) + assert probe.releases == 0 + assert monitor.processed == 3 + assert monitor.successes == 0 + assert not monitor.pending + probe.found = frozenset(f"response-{i}" for i in range(3, 303)) + assert not monitor.poll(tuple(record(i) for i in range(3, 303))) + monitor.assert_complete(state("complete", pending=False)) + assert probe.releases == 1 + assert monitor.processed == 303 + assert monitor.successes == monitor.persisted_count == 300 + + +def test_fourth_non_200_fails_across_poll_batches() -> None: + monitor: Final = Probe().monitor() + assert not monitor.poll(tuple(failed_record(i) for i in range(3))) + with pytest.raises(AssertionError, match="Too many unexpected non-200"): + _ = monitor.poll((failed_record(3),)) + + +@pytest.mark.parametrize( + "failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})] +) +def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: + with pytest.raises(AssertionError, match="Failed or empty inference"): + _ = Probe().monitor().poll((failed,)) + + +@pytest.mark.parametrize( + "failures,overhead,ttft,growth,expected", + [ + (3, 0.29, 3.7, 100, ()), + (4, 0.29, 3.7, 100, ("unexpected_status_count",)), + (3, 0.30, 3.7, 100, ("gateway_added_seconds percentile 50",)), + (3, 0.29, 3.8, 100, ("time_to_first_token_seconds percentile 99",)), + (3, 0.29, 3.7, 101, ("memory_growth_megabytes",)), + ], +) +def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples( + failures: int, overhead: float, ttft: float, growth: float, expected: tuple[str, ...] +) -> None: + nonstream: Final = tuple( + record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100) + ) + stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft}) + observed: Final = Observation( + records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))), + mock_requests=(), + guardrail_calls=(), + rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()), + rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()), + wall_seconds=60, + ) + rows, exceeded = evaluate(upgrade_case(), observed) + assert tuple(row.label for row in rows if row.over or row.missing) == expected + assert len(exceeded) == len(expected) diff --git a/tests/integration/_support/bedrock_realtime_peer.py b/tests/integration/_support/bedrock_realtime_peer.py index d30bf1b92f5..ff618a00136 100644 --- a/tests/integration/_support/bedrock_realtime_peer.py +++ b/tests/integration/_support/bedrock_realtime_peer.py @@ -14,6 +14,12 @@ text message (the injected trigger-audio block is USER/AUDIO, so it never trips ``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration`` ``toolUseId`` names the call being answered. + +A ``Refusal`` peer withholds each stream's initial response until that stream's opening arrived +(its SYSTEM prompt block or its first USER text block closed), then answers the HTTP error Bedrock +sends for an unknown model id. A ``StreamFailure`` peer accepts the stream, plays its turns, and +then sends an event-stream exception frame and ends the stream, as Bedrock does for a mid-call +model error. ``streams()`` groups the inbound events per opened stream. """ from __future__ import annotations @@ -26,22 +32,24 @@ import uuid from collections.abc import Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass -from itertools import chain +from itertools import chain, count from pathlib import Path from queue import SimpleQueue -from typing import Final, Literal +from typing import Final, Literal, TypeVar import h2.config import h2.connection import h2.events import h2.exceptions from integration._support.async_server import Handler, LoopbackServer -from integration._support.process import _free_port +from integration._support.ports import allocate_port from integration._support.tls import server_context from integration._support.upstream import aws_event_stream_frame from pydantic import JsonValue, TypeAdapter BedrockEvent = Mapping[str, JsonValue] +OpeningEnd = Literal["system_prompt", "user_input"] +Recorded = TypeVar("Recorded") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -177,6 +185,20 @@ class Turn: events: tuple[BedrockEvent, ...] +@dataclass(frozen=True, slots=True) +class Refusal: + after: OpeningEnd + status: int = 400 + error_type: str = "ValidationException" + message: str = "The provided model identifier is invalid." + + +@dataclass(frozen=True, slots=True) +class StreamFailure: + exception_type: str = "modelStreamErrorException" + message: str = "The system encountered an unexpected error during processing. Try your request again." + + def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]: body: Final = event["event"] assert isinstance(body, Mapping) and len(body) == 1, event @@ -214,6 +236,17 @@ def _frame_payload(raw: bytes) -> bytes: return raw[12 + header_length : -4] +def _exception_frame(failure: StreamFailure) -> bytes: + return aws_event_stream_frame( + { + ":message-type": "exception", + ":exception-type": failure.exception_type, + ":content-type": "application/json", + }, + json.dumps({"message": failure.message}).encode(), + ) + + def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None: """One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event.""" inner_payload: Final = _frame_payload(chunk_payload) @@ -239,6 +272,7 @@ class _Inbound: self.tool_use_ids: Final[dict[str, str]] = {} self.tool_results_seen = 0 self.user_text_closed = False + self.system_text_closed = False def record(self, event: BedrockEvent) -> None: self.events.put(event) @@ -253,6 +287,8 @@ class _Inbound: role_type: Final = self.open_blocks.pop(content_name, ("", "")) if role_type == ("USER", "TEXT"): self.user_text_closed = True + elif role_type == ("SYSTEM", "TEXT"): + self.system_text_closed = True elif name == "toolResult": self.tool_results_seen += 1 @@ -283,10 +319,35 @@ class _Script: return turn.events +def _opening_arrived(after: OpeningEnd, inbound: _Inbound) -> bool: + match after: + case "system_prompt": + return inbound.system_text_closed + case "user_input": + return inbound.user_text_closed + + class _PeerState: - def __init__(self, turns: tuple[Turn, ...]) -> None: + def __init__(self, turns: tuple[Turn, ...], refusal: Refusal | None, failure: StreamFailure | None) -> None: self.inbound: Final = _Inbound() self.script: Final = _Script(turns) + self.refusal: Final = refusal + self.failure: Final = failure + self.ordinals: Final = count() + self.opened: Final = SimpleQueue[int]() + self.stream_events: Final = SimpleQueue[tuple[int, BedrockEvent]]() + + +def _drain(queue: SimpleQueue[Recorded], log: list[Recorded]) -> tuple[Recorded, ...]: + while not queue.empty(): + log.append(queue.get()) # drained log, not a built value + return tuple(log) + + +@dataclass(frozen=True, slots=True) +class _Stream: + ordinal: int + inbound: _Inbound @dataclass(frozen=True, slots=True) @@ -294,11 +355,16 @@ class BedrockRealtimePeer: api_base: str _state: _PeerState _drained: list[BedrockEvent] + _opened: list[int] + _stream_events: list[tuple[int, BedrockEvent]] def received(self) -> tuple[BedrockEvent, ...]: - while not self._state.inbound.events.empty(): - self._drained.append(self._state.inbound.events.get()) # drained log, not a built value - return tuple(self._drained) + return _drain(self._state.inbound.events, self._drained) + + def streams(self) -> tuple[tuple[BedrockEvent, ...], ...]: + opened: Final = _drain(self._state.opened, self._opened) + recorded: Final = _drain(self._state.stream_events, self._stream_events) + return tuple(tuple(event for ordinal, event in recorded if ordinal == stream) for stream in opened) def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]: return tuple(chain.from_iterable(_tool_specs(event) for event in self.received())) @@ -344,6 +410,7 @@ class _H2Session: self._writer: Final = writer self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False)) self._buffers: Final[dict[int, bytes]] = {} + self._streams: Final[dict[int, _Stream]] = {} self._ended: set[int] = set() self._closed = False @@ -367,10 +434,14 @@ class _H2Session: def _apply(self, event: h2.events.Event) -> None: if isinstance(event, h2.events.RequestReceived): - self._h2.send_headers( - event.stream_id, - [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], - ) + ordinal: Final = next(self._state.ordinals) + self._streams[event.stream_id] = _Stream(ordinal, _Inbound()) + self._state.opened.put(ordinal) + if self._state.refusal is None: + self._h2.send_headers( + event.stream_id, + [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], + ) elif isinstance(event, h2.events.DataReceived): self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id) self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data @@ -380,15 +451,44 @@ class _H2Session: decoded: Final = _decode_inbound(payload) if decoded is None: continue - self._state.inbound.record(decoded) - if event.stream_id not in self._ended: - for outbound in self._state.script.ready_events(self._state.inbound, decoded): - self._h2.send_data(event.stream_id, _outbound_chunk(outbound)) + self._record(event.stream_id, decoded) elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)): self._close_stream(event.stream_id) elif isinstance(event, h2.events.ConnectionTerminated): self._writer.close() + def _record(self, stream_id: int, decoded: BedrockEvent) -> None: + stream: Final = self._streams[stream_id] + self._state.inbound.record(decoded) + stream.inbound.record(decoded) + self._state.stream_events.put((stream.ordinal, decoded)) + if stream_id in self._ended: + return + refusal: Final = self._state.refusal + if refusal is not None: + if _opening_arrived(refusal.after, stream.inbound): + self._refuse(stream_id, refusal) + return + outbound_events: Final = self._state.script.ready_events(self._state.inbound, decoded) + for outbound in outbound_events: + self._h2.send_data(stream_id, _outbound_chunk(outbound)) + failure: Final = self._state.failure + if outbound_events and failure is not None and self._state.script.exhausted: + self._h2.send_data(stream_id, _exception_frame(failure)) + self._end_stream(stream_id) + + def _refuse(self, stream_id: int, refusal: Refusal) -> None: + self._h2.send_headers( + stream_id, + [ + (":status", str(refusal.status)), + ("content-type", "application/json"), + ("x-amzn-errortype", refusal.error_type), + ], + ) + self._h2.send_data(stream_id, json.dumps({"message": refusal.message}).encode()) + self._end_stream(stream_id) + def _close_stream(self, stream_id: int) -> None: self._end_stream(stream_id) self._closed = True @@ -411,19 +511,27 @@ def _session_handler(state: _PeerState) -> Handler: @contextmanager -def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]: +def bedrock_realtime_peer( + turns: tuple[Turn, ...], + cert: tuple[Path, Path], + *, + refusal: Refusal | None = None, + failure: StreamFailure | None = None, +) -> Iterator[BedrockRealtimePeer]: """Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer. `cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE. SigV4 headers are ignored; the stream closes cleanly when the client ends it. """ - port: Final = _free_port() + port: Final = allocate_port() context: Final = server_context(cert[0], cert[1]) context.set_alpn_protocols(["h2"]) - state: Final = _PeerState(turns) + state: Final = _PeerState(turns, refusal, failure) server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context) server.start() try: - yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[]) + yield BedrockRealtimePeer( + api_base=f"https://localhost:{port}", _state=state, _drained=[], _opened=[], _stream_events=[] + ) finally: server.stop() diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index fd61ca3aea2..00b9dca5184 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -628,8 +628,11 @@ def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}" contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text()) - directory.mkdir(parents=True, exist_ok=True) - destination: Final = directory / f"worker-config-{uuid.uuid4().hex}.yaml" + checkouts: Final = (_harness_root(), _proxy_root().resolve()) + inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts) + target: Final = directory if inside_checkout else config.parent + target.mkdir(parents=True, exist_ok=True) + destination: Final = target / f"worker-config-{uuid.uuid4().hex}.yaml" loaded: Final[object] = yaml.safe_load(contents) if not isinstance(loaded, dict): destination.write_text(contents) diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index e5ce909ab18..df6036ff8d7 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -59,6 +59,7 @@ class Reply: content_type: str = "application/json" chunks: tuple[bytes, ...] | None = None abort_after: int | None = None + gate_before_first: threading.Event | None = None gate_after_first: threading.Event | None = None pause_between_chunks: float = 0 heartbeat: bytes | None = None @@ -166,6 +167,12 @@ def wire_server( for index, chunk in enumerate(reply.chunks): if reply.abort_after == index: break + if ( + index == 0 + and reply.gate_before_first is not None + and not _await_release(reply.gate_before_first, closing) + ): + break self.write_chunk(chunk) if ( index == 0 diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index 42ee063ddc7..dd98ff6dea1 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -1695,15 +1695,19 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( record_property("chaos_spend_row_counts", str(tuple(len(rows) for rows in spend_rows))) with owned.replacing_worker(min(workers)) as kill: kill() - post_kill: Final = tuple( - _post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5] - ) - assert all( - response.status_code == 400 and response.json() == expected - for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) - ), [response.text for response in post_kill] - recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)) - assert recovered.status_code == 200, recovered.text + with httpx.Client(base_url=candidate.client.base_url, timeout=15, trust_env=False) as fresh: + survivor: Final = Gateway(fresh, candidate.key, candidate.upstream_url) + post_kill: Final = tuple( + _post(survivor, path, body) for path, body, _expected, _identity in missing_calls[:5] + ) + assert all( + response.status_code == 400 and response.json() == expected + for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) + ), [response.text for response in post_kill] + recovered: Final = _post( + survivor, "/v1/chat/completions", _chat_body(chat_model, "recovered", True) + ) + assert recovered.status_code == 200, recovered.text def test_eval_create_and_run_forward_client_metadata(gateway: Gateway) -> None: diff --git a/tests/integration/configuration/test_effective_settings.py b/tests/integration/configuration/test_effective_settings.py index 7f6313e209e..014cfd3a866 100644 --- a/tests/integration/configuration/test_effective_settings.py +++ b/tests/integration/configuration/test_effective_settings.py @@ -1,5 +1,7 @@ import json import os +import subprocess +import sys import uuid from collections.abc import Callable, Iterator from contextlib import contextmanager @@ -7,9 +9,13 @@ from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Final +from urllib.parse import parse_qs, urlsplit, urlunsplit import httpx +import psycopg import pytest +from psycopg import sql +from psycopg.types.json import Jsonb from pydantic import BaseModel, JsonValue, TypeAdapter from tests.integration._support.client import Gateway, eventually, object_value, string_value @@ -1111,3 +1117,585 @@ def test_credential_value_update_and_model_reload_reach_provider(gateway: Gatewa observed.raise_for_status() assert len(observed.json()["requests"]) == 1, (value, reload, observed.text) assert observed.json()["requests"][0]["authorization"] == f"Bearer {value}" + + +_ENVIRONMENT_VARIABLES_REVISION_MIGRATION: Final = ( + "litellm-proxy-extras/litellm_proxy_extras/migrations/" + "20261015120000_add_environment_variables_revision_trigger/migration.sql" +) + + +def _prisma(*arguments: str) -> None: + subprocess.run((sys.executable, "-I", "-m", "prisma", *arguments), check=True, capture_output=True, text=True) + + +def _push_schema() -> None: + _prisma("db", "push", "--schema", "litellm/proxy/schema.prisma", "--skip-generate") + + +def _push_schema_with_environment_variables_revision_trigger() -> None: + _push_schema() + _prisma( + "db", "execute", "--schema", "litellm/proxy/schema.prisma", "--file", _ENVIRONMENT_VARIABLES_REVISION_MIGRATION + ) + + +def _provider_echoing_its_key(source: str) -> Callable[[Request], Reply]: + def provider(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/chat/completions", request.target + return _completion_reply(f"{source} {request.headers['authorization']}") + + return provider + + +def _environment_variables_proxy_config( + directory: Path, gateway: Gateway, alias: str, config_file_provider: Wire +) -> Path: + config: Final = directory / "environment-variables-proxy.yaml" + config.write_text( + json.dumps( + { + "model_list": [{"model_name": alias, "litellm_params": {"model": "openai/gpt-4o-mini"}}], + "general_settings": {"master_key": gateway.key, "store_model_in_db": True}, + "environment_variables": {"OPENAI_BASE_URL": config_file_provider.url + "/v1"}, + } + ) + ) + return config + + +def _environment_variables_proxy_environment(database_url: str, process_key: str) -> dict[str, str]: + return { + "DATABASE_URL": database_url, + "DISABLE_SCHEMA_UPDATE": "true", + "LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS": "1", + "OPENAI_API_KEY": process_key, + "PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2", + } + + +def _chat_answer(proxy: Gateway, alias: str) -> str: + response: Final = proxy.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": [{"role": "user", "content": "which key"}]} + ) + assert response.status_code == 200, response.text + return string_value(object_value(object_value(response.json()["choices"][0])["message"])["content"]) + + +def test_removing_db_environment_variables_restores_process_and_config_file_values_on_every_proxy( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + alias: Final = "environment-variables-" + marker + process_key: Final = "process-key-" + marker + database_key: Final = "database-key-" + marker + + with ( + scratch_database() as database_url, + wire_server(_provider_echoing_its_key("config-file")) as config_file_provider, + wire_server(_provider_echoing_its_key("database")) as database_provider, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + _push_schema_with_environment_variables_revision_trigger() + config: Final = _environment_variables_proxy_config(tmp_path, gateway, alias, config_file_provider) + environment: Final = _environment_variables_proxy_environment(database_url, process_key) + removed: Final = ("DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE", "OPENAI_BASE_URL") + from_process_and_config_file: Final = f"config-file Bearer {process_key}" + from_database: Final = f"database Bearer {database_key}" + with owned_proxy( + gateway, tmp_path / "first", environment, config=config, database_setup=(), remove_environment=removed + ) as first: + assert _chat_answer(first, alias) == from_process_and_config_file + + stored: Final = first.request( + "POST", + "/config/update", + { + "environment_variables": { + "OPENAI_API_KEY": database_key, + "OPENAI_BASE_URL": database_provider.url + "/v1", + } + }, + ) + assert stored.status_code == 200, stored.text + eventually(lambda: _chat_answer(first, alias), lambda value: value == from_database, seconds=30) + + with owned_proxy( + gateway, tmp_path / "peer", environment, config=config, database_setup=(), remove_environment=removed + ) as peer_booted_with_the_row_stored: + + def answers() -> tuple[str, str]: + return (_chat_answer(first, alias), _chat_answer(peer_booted_with_the_row_stored, alias)) + + eventually(answers, lambda values: values == (from_database,) * 2, seconds=30) + + unset: Final = first.request( + "POST", + "/config/update", + {"environment_variables": {"OPENAI_API_KEY": None, "OPENAI_BASE_URL": None}}, + ) + assert unset.status_code == 200, unset.text + assert _config_section(database_url, "environment_variables") == { + "OPENAI_API_KEY": None, + "OPENAI_BASE_URL": None, + } + eventually(answers, lambda values: values == (from_process_and_config_file,) * 2, seconds=30) + + +def test_a_proxy_that_pushes_its_own_schema_unsets_a_removed_db_environment_variable( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + alias: Final = "environment-variables-db-push-" + marker + process_key: Final = "process-key-" + marker + database_key: Final = "database-key-" + marker + + with ( + scratch_database() as database_url, + wire_server(_provider_echoing_its_key("config-file")) as config_file_provider, + wire_server(_provider_echoing_its_key("database")) as database_provider, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + config: Final = _environment_variables_proxy_config(tmp_path, gateway, alias, config_file_provider) + environment: Final = { + **_environment_variables_proxy_environment(database_url, process_key), + "DISABLE_SCHEMA_UPDATE": "false", + } + from_process_and_config_file: Final = f"config-file Bearer {process_key}" + with owned_proxy( + gateway, + tmp_path / "proxy", + environment, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE", "OPENAI_BASE_URL"), + ) as proxy: + assert _chat_answer(proxy, alias) == from_process_and_config_file + stored: Final = proxy.request( + "POST", + "/config/update", + { + "environment_variables": { + "OPENAI_API_KEY": database_key, + "OPENAI_BASE_URL": database_provider.url + "/v1", + } + }, + ) + assert stored.status_code == 200, stored.text + eventually( + lambda: _chat_answer(proxy, alias), lambda value: value == f"database Bearer {database_key}", seconds=30 + ) + + unset: Final = proxy.request( + "POST", + "/config/update", + {"environment_variables": {"OPENAI_API_KEY": None, "OPENAI_BASE_URL": None}}, + ) + assert unset.status_code == 200, unset.text + eventually( + lambda: _chat_answer(proxy, alias), lambda value: value == from_process_and_config_file, seconds=30 + ) + + +def _environment_variables_row(database_url: str) -> dict[str, JsonValue]: + return read_rows( + 'SELECT param_value, reload_revision FROM "LiteLLM_Config" WHERE param_name = %s', + ("environment_variables",), + database_url=database_url, + )[0] + + +def _load_environment_variables_row(replica_url: str, row: dict[str, JsonValue]) -> None: + with psycopg.connect(replica_url, autocommit=True) as replica: + replica.execute( + 'INSERT INTO "LiteLLM_Config" (param_name, param_value, reload_revision) VALUES (%s, %s, %s) ' + "ON CONFLICT (param_name) DO UPDATE " + "SET param_value = EXCLUDED.param_value, reload_revision = EXCLUDED.reload_revision", + ("environment_variables", Jsonb(row["param_value"]), row["reload_revision"]), + ) + + +def test_a_read_replica_back_on_an_older_row_does_not_bring_back_a_removed_db_environment_variable( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + alias: Final = "environment-variables-replica-" + marker + process_key: Final = "process-key-" + marker + database_key: Final = "database-key-" + marker + + with ( + scratch_database() as writer_url, + scratch_database() as replica_url, + wire_server(_provider_echoing_its_key("config-file")) as config_file_provider, + wire_server(_provider_echoing_its_key("database")) as database_provider, + wire_server(_provider_echoing_its_key("writer")) as writer_provider, + ): + monkeypatch.setenv("DATABASE_URL", replica_url) + _push_schema() + monkeypatch.setenv("DATABASE_URL", writer_url) + _push_schema_with_environment_variables_revision_trigger() + config: Final = _environment_variables_proxy_config(tmp_path, gateway, alias, config_file_provider) + environment: Final = { + **_environment_variables_proxy_environment(writer_url, process_key), + "DATABASE_URL_READ_REPLICA": replica_url, + } + from_process_and_config_file: Final = f"config-file Bearer {process_key}" + with owned_proxy( + gateway, + tmp_path / "replica-reader", + environment, + config=config, + database_setup=(), + remove_environment=("OPENAI_API_BASE", "OPENAI_BASE_URL"), + ) as proxy: + assert _chat_answer(proxy, alias) == from_process_and_config_file + + stored: Final = proxy.request( + "POST", + "/config/update", + { + "environment_variables": { + "OPENAI_API_KEY": database_key, + "OPENAI_BASE_URL": database_provider.url + "/v1", + } + }, + ) + assert stored.status_code == 200, stored.text + row_holding_the_key: Final = _environment_variables_row(writer_url) + _load_environment_variables_row(replica_url, row_holding_the_key) + eventually( + lambda: _chat_answer(proxy, alias), + lambda value: value == f"database Bearer {database_key}", + seconds=30, + ) + + unset: Final = proxy.request( + "POST", "/config/update", {"environment_variables": {"OPENAI_API_KEY": None, "OPENAI_BASE_URL": None}} + ) + assert unset.status_code == 200, unset.text + _load_environment_variables_row(replica_url, _environment_variables_row(writer_url)) + eventually( + lambda: _chat_answer(proxy, alias), lambda value: value == from_process_and_config_file, seconds=30 + ) + + _load_environment_variables_row(replica_url, row_holding_the_key) + _save_environment_variables_as(writer_url, {"OPENAI_BASE_URL": writer_provider.url + "/v1"}) + eventually( + lambda: _chat_answer(proxy, alias), + lambda value: value == f"writer Bearer {process_key}", + seconds=30, + ) + + +_CAPPED_WRITER_ROLE_PASSWORD: Final = "integration-capped-writer-password" + + +@contextmanager +def _database_owned_by_a_login_role() -> Iterator[tuple[str, str, str]]: + role: Final = f"integration_capped_writer_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD {} NOSUPERUSER NOINHERIT").format( + sql.Identifier(role), sql.Literal(_CAPPED_WRITER_ROLE_PASSWORD) + ) + ) + try: + with scratch_database() as admin_url: + parsed: Final = urlsplit(admin_url) + admin.execute( + sql.SQL("ALTER DATABASE {} OWNER TO {}").format( + sql.Identifier(parsed.path.lstrip("/")), sql.Identifier(role) + ) + ) + host: Final = parsed.netloc.rpartition("@")[2] + yield ( + role, + admin_url, + urlunsplit(parsed._replace(netloc=f"{role}:{_CAPPED_WRITER_ROLE_PASSWORD}@{host}")), + ) + finally: + admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) + + +def _set_connection_limit(admin_url: str, role: str, limit: int) -> None: + with psycopg.connect(admin_url, autocommit=True) as admin: + admin.execute(sql.SQL("ALTER ROLE {} CONNECTION LIMIT {}").format(sql.Identifier(role), sql.Literal(limit))) + + +def _open_connections(database_url: str, role: str) -> JsonValue: + return read_rows( + "SELECT count(*) AS connections FROM pg_stat_activity WHERE usename = %s", + (role,), + database_url=database_url, + )[0]["connections"] + + +def test_a_writer_database_out_of_connections_keeps_the_applied_db_environment_variables( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + alias: Final = "environment-variables-capped-writer-" + marker + process_key: Final = "process-key-" + marker + database_key: Final = "database-key-" + marker + + with ( + _database_owned_by_a_login_role() as (role, writer_admin_url, writer_url), + scratch_database() as replica_url, + wire_server(_provider_echoing_its_key("config-file")) as config_file_provider, + wire_server(_provider_echoing_its_key("database")) as database_provider, + wire_server(_provider_echoing_its_key("writer")) as writer_provider, + ): + monkeypatch.setenv("DATABASE_URL", replica_url) + _push_schema() + monkeypatch.setenv("DATABASE_URL", writer_url) + _push_schema_with_environment_variables_revision_trigger() + config: Final = _environment_variables_proxy_config(tmp_path, gateway, alias, config_file_provider) + parsed_writer_url: Final = urlsplit(writer_url) + writer_url_dropping_idle_connections: Final = urlunsplit( + parsed_writer_url._replace( + query="&".join(part for part in (parsed_writer_url.query, "max_idle_connection_lifetime=1") if part) + ) + ) + environment: Final = { + **_environment_variables_proxy_environment(writer_url_dropping_idle_connections, process_key), + "DATABASE_URL_READ_REPLICA": replica_url, + } + from_process_and_config_file: Final = f"config-file Bearer {process_key}" + with owned_proxy( + gateway, + tmp_path / "capped-writer", + environment, + config=config, + database_setup=(), + remove_environment=("OPENAI_API_BASE", "OPENAI_BASE_URL"), + ) as proxy: + stored: Final = proxy.request( + "POST", + "/config/update", + { + "environment_variables": { + "OPENAI_API_KEY": database_key, + "OPENAI_BASE_URL": database_provider.url + "/v1", + } + }, + ) + assert stored.status_code == 200, stored.text + row_holding_the_key: Final = _environment_variables_row(writer_admin_url) + _load_environment_variables_row(replica_url, row_holding_the_key) + eventually( + lambda: _chat_answer(proxy, alias), + lambda value: value == f"database Bearer {database_key}", + seconds=30, + ) + + unset: Final = proxy.request( + "POST", "/config/update", {"environment_variables": {"OPENAI_API_KEY": None, "OPENAI_BASE_URL": None}} + ) + assert unset.status_code == 200, unset.text + _load_environment_variables_row(replica_url, _environment_variables_row(writer_admin_url)) + eventually( + lambda: _chat_answer(proxy, alias), lambda value: value == from_process_and_config_file, seconds=30 + ) + + _load_environment_variables_row(replica_url, row_holding_the_key) + _set_connection_limit(writer_admin_url, role, 0) + with pytest.raises(psycopg.OperationalError): + psycopg.connect(writer_url).close() + eventually(lambda: _open_connections(writer_admin_url, role), lambda count: count == 0, seconds=30) + + config_read: Final = proxy.request("GET", "/get/config/callbacks") + assert config_read.status_code == 200, config_read.text + assert _chat_answer(proxy, alias) == from_process_and_config_file + + _set_connection_limit(writer_admin_url, role, -1) + _save_environment_variables_as(writer_admin_url, {"OPENAI_BASE_URL": writer_provider.url + "/v1"}) + eventually( + lambda: _chat_answer(proxy, alias), + lambda value: value == f"writer Bearer {process_key}", + seconds=30, + ) + + +def _sign_in_client_id(proxy: Gateway) -> str | None: + response: Final = httpx.get( + str(proxy.client.base_url).rstrip("/") + "/sso/key/generate", + follow_redirects=False, + trust_env=False, + timeout=30, + ) + if not response.is_redirect: + assert response.status_code == 200, response.text + return None + return parse_qs(urlsplit(response.headers["location"]).query)["client_id"][0] + + +def test_an_sso_client_removed_from_db_environment_variables_stops_sign_in_on_every_proxy( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + environment_row_client: Final = "environment-row-client-" + marker + sso_table_client: Final = "sso-table-client-" + marker + + with scratch_database() as database_url, wire_server(_provider_echoing_its_key("config-file")) as provider: + monkeypatch.setenv("DATABASE_URL", database_url) + _push_schema_with_environment_variables_revision_trigger() + config: Final = _environment_variables_proxy_config(tmp_path, gateway, "sso-" + marker, provider) + environment: Final = _environment_variables_proxy_environment(database_url, "process-key-" + marker) + removed: Final = ( + "DATABASE_URL_READ_REPLICA", + "GENERIC_CLIENT_ID", + "GOOGLE_CLIENT_ID", + "GOOGLE_CLIENT_SECRET", + "MICROSOFT_CLIENT_ID", + "PROXY_BASE_URL", + ) + with ( + owned_proxy( + gateway, tmp_path / "first", environment, config=config, database_setup=(), remove_environment=removed + ) as first, + owned_proxy( + gateway, tmp_path / "peer", environment, config=config, database_setup=(), remove_environment=removed + ) as peer, + ): + + def sign_in_clients() -> tuple[str | None, str | None]: + return (_sign_in_client_id(first), _sign_in_client_id(peer)) + + def store_environment_row_client(client_id: str | None, secret: str | None) -> None: + response: Final = first.request( + "POST", + "/config/update", + {"environment_variables": {"GOOGLE_CLIENT_ID": client_id, "GOOGLE_CLIENT_SECRET": secret}}, + ) + assert response.status_code == 200, response.text + + assert sign_in_clients() == (None, None) + + store_environment_row_client(environment_row_client, "environment-row-secret") + eventually(sign_in_clients, lambda clients: clients == (environment_row_client,) * 2, seconds=30) + + store_environment_row_client(None, None) + eventually(sign_in_clients, lambda clients: clients == (None, None), seconds=30) + + store_environment_row_client(environment_row_client, "environment-row-secret") + eventually(sign_in_clients, lambda clients: clients == (environment_row_client,) * 2, seconds=30) + moved: Final = first.request( + "PATCH", + "/update/sso_settings", + {"google_client_id": sso_table_client, "google_client_secret": "sso-table-secret"}, + ) + assert moved.status_code == 200, moved.text + eventually(sign_in_clients, lambda clients: clients == (sso_table_client,) * 2, seconds=30) + + +_TABLE_ONLY_ROLE_PASSWORD: Final = "integration-table-only-password" + + +@contextmanager +def _login_role_with_no_sequence_grants() -> Iterator[str]: + role: Final = f"integration_table_only_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD {} NOSUPERUSER NOINHERIT").format( + sql.Identifier(role), sql.Literal(_TABLE_ONLY_ROLE_PASSWORD) + ) + ) + try: + yield role + finally: + admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) + + +def _save_environment_variables_as(database_url: str, values: dict[str, str]) -> None: + with psycopg.connect(database_url, autocommit=True) as writer: + writer.execute( + 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s) ' + "ON CONFLICT (param_name) DO UPDATE SET param_value = EXCLUDED.param_value", + ("environment_variables", Jsonb(values)), + ) + + +def _environment_variables_revision(database_url: str) -> JsonValue: + return read_rows( + 'SELECT reload_revision FROM "LiteLLM_Config" WHERE param_name = %s', + ("environment_variables",), + database_url=database_url, + )[0]["reload_revision"] + + +def test_a_database_role_with_no_sequence_grants_can_still_save_db_environment_variables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with _login_role_with_no_sequence_grants() as role, scratch_database() as database_url: + monkeypatch.setenv("DATABASE_URL", database_url) + _push_schema_with_environment_variables_revision_trigger() + with psycopg.connect(database_url, autocommit=True) as owner: + owner.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role))) + owner.execute( + sql.SQL("GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA public TO {}").format( + sql.Identifier(role) + ) + ) + parsed: Final = urlsplit(database_url) + role_url: Final = urlunsplit( + parsed._replace(netloc=f"{role}:{_TABLE_ONLY_ROLE_PASSWORD}@{parsed.netloc.rpartition('@')[2]}") + ) + + _save_environment_variables_as(role_url, {"TABLE_ONLY_ROLE_KEY": "stored"}) + first: Final = _environment_variables_revision(database_url) + _save_environment_variables_as(role_url, {}) + second: Final = _environment_variables_revision(database_url) + + assert isinstance(first, int) and isinstance(second, int) + assert 0 < first < second + + +def test_db_environment_variables_revision_keeps_rising_after_its_sequence_falls_behind_the_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with scratch_database() as database_url: + monkeypatch.setenv("DATABASE_URL", database_url) + _push_schema_with_environment_variables_revision_trigger() + for value in ("first", "second", "third"): + _save_environment_variables_as(database_url, {"SEQUENCE_BEHIND_KEY": value}) + before_the_sequence_restarts: Final = _environment_variables_revision(database_url) + with psycopg.connect(database_url, autocommit=True) as owner: + owner.execute('ALTER SEQUENCE "litellm_environment_variables_revision_seq" RESTART WITH 1') + + _save_environment_variables_as(database_url, {}) + after_the_sequence_restarts: Final = _environment_variables_revision(database_url) + + assert isinstance(before_the_sequence_restarts, int) and isinstance(after_the_sequence_restarts, int) + assert after_the_sequence_restarts > before_the_sequence_restarts + + +def test_db_environment_variables_saved_from_a_session_outside_the_tables_schema_still_get_a_revision( + monkeypatch: pytest.MonkeyPatch, +) -> None: + schema: Final = "litellm_tenant" + table: Final = sql.SQL('{}."LiteLLM_Config"').format(sql.Identifier(schema)) + with scratch_database() as database_url: + monkeypatch.setenv("DATABASE_URL", f"{database_url}?schema={schema}") + _push_schema_with_environment_variables_revision_trigger() + + def save_and_read_revision(values: dict[str, str]) -> object: + with psycopg.connect(database_url, autocommit=True) as writer: + writer.execute( + sql.SQL( + "INSERT INTO {} (param_name, param_value) VALUES (%s, %s) " + "ON CONFLICT (param_name) DO UPDATE SET param_value = EXCLUDED.param_value" + ).format(table), + ("environment_variables", Jsonb(values)), + ) + row: Final = writer.execute( + sql.SQL("SELECT reload_revision FROM {} WHERE param_name = %s").format(table), + ("environment_variables",), + ).fetchone() + assert row is not None + return row[0] + + first: Final = save_and_read_revision({"TENANT_SCHEMA_KEY": "stored"}) + second: Final = save_and_read_revision({}) + + assert isinstance(first, int) and isinstance(second, int) + assert 0 < first < second diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index a0c81d8f06d..5ad4afa9530 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -6899,6 +6899,35 @@ CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_T } ) +def _responses_bridged_stream_case( + sibling: CostTrackingTestCase, stream_sibling: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "request": {**sibling.request, "stream": True}, + "response": stream_sibling.response, + } + ) + + +CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + "claude-fable-5-1-responses_bridged_stream", +) +CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + "claude-haiku-5-5-responses_bridged_stream", +) +CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_MESSAGES_STREAM, + "claude-opus-5-5-responses_bridged_stream", +) +CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_MESSAGES_STREAM, + "claude-sonnet-5-5-responses_bridged_stream", +) + CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( CLAUDE_HAIKU_5_5_INPUT_TEXT, "claude-haiku-5-5-tiered_input_above_100k", @@ -6994,6 +7023,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, @@ -7119,6 +7152,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( + StreamParityTestCase( + plain=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM + ), StreamParityTestCase(plain=CLAUDE_HAIKU_4_5_INPUT_TEXT, streamed=CLAUDE_HAIKU_4_5_STREAM), StreamParityTestCase( plain=CLAUDE_HAIKU_4_5_RESPONSE_MODEL_OVERRIDE, streamed=CLAUDE_HAIKU_4_5_STREAM_RESPONSE_MODEL_OVERRIDE diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 1360f65b802..d9f17c12bda 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -21,6 +21,8 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.cases.openai import _responses_json_body, _responses_stream_case +from pydantic import JsonValue def _response_for_model( @@ -47,6 +49,37 @@ def _stream_for_model(response: SseResponse, model: str) -> SseResponse: ) +def _azure_responses_json_case( + base: CostTrackingTestCase, name: str, response_model: str, content_filters: JsonValue +) -> CostTrackingTestCase: + response_case: Final = bridge_pairs((base,), ("/v1/responses",))[0].bridged + return response_case.model_copy( + update={ + "name": name, + "response": JsonResponse( + content_type="application/json", + body=_responses_json_body( + response_model, + include_content_filters=True, + content_filters=content_filters, + ), + ), + } + ) + + +def _azure_messages_stream_case( + messages_json: CostTrackingTestCase, stream_case: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return messages_json.model_copy( + update={ + "name": name, + "request": {**messages_json.request, "stream": True}, + "response": stream_case.response, + } + ) + + # azure/gpt-5.4-mini AZURE_GPT_5_4_MINI_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-5.4-mini-input_text", @@ -5812,6 +5845,85 @@ AZURE_GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( ExactExpected(spend=0.22375, input_cost=0.2, output_cost=0.02375, prompt_tokens=12928, completion_tokens=380), ) +AZURE_GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + "azure/gpt-5.6-terra-tiered_input_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + }, + ExactExpected( + spend=1.3281576, input_cost=1.32, output_cost=0.0081576, prompt_tokens=300000, completion_tokens=412 + ), +) +AZURE_GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K: Final = with_usage( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + "azure/gpt-5.6-terra-tiered_cache_read_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, + }, + ExactExpected( + spend=0.5361576, input_cost=0.528, output_cost=0.0081576, prompt_tokens=300000, completion_tokens=412 + ), +) + + +AZURE_GPT_6_1_SOL_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_1_SOL_INPUT_TEXT, "azure-gpt-6.1-sol-responses_json", "gpt-6.1-sol", None +) +AZURE_GPT_6_LUNA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_LUNA_INPUT_TEXT, "azure-gpt-6-luna-responses_json", "gpt-6-luna", None +) +AZURE_GPT_6_ASTRA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_ASTRA_INPUT_TEXT, "azure-gpt-6-astra-responses_json", "gpt-6-astra", [] +) +AZURE_GPT_5_6_TERRA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, "azure-gpt-5.6-terra-responses_json", "gpt-5.6-terra", None +) +AZURE_GPT_6_1_SOL_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_1_SOL_RESPONSES_JSON, + "azure-gpt-6.1-sol-responses_stream", + "gpt-6.1-sol", + 14, + include_content_filters=True, +) +AZURE_GPT_6_LUNA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_LUNA_RESPONSES_JSON, + "azure-gpt-6-luna-responses_stream", + "gpt-6-luna", + 25, + include_content_filters=True, +) +AZURE_GPT_6_ASTRA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_ASTRA_RESPONSES_JSON, + "azure-gpt-6-astra-responses_stream", + "gpt-6-astra", + 16, + include_content_filters=True, + content_filters=[], +) +AZURE_GPT_5_6_TERRA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_5_6_TERRA_RESPONSES_JSON, + "azure-gpt-5.6-terra-responses_stream", + "gpt-5.6-terra", + 15, + include_content_filters=True, +) +AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_1_SOL_STREAM, "azure-gpt-6.1-sol-messages_bridged_stream" +) +AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_LUNA_STREAM, "azure-gpt-6-luna-messages_bridged_stream" +) +AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_ASTRA_STREAM, "azure-gpt-6-astra-messages_bridged_stream" +) +AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, AZURE_GPT_5_6_TERRA_STREAM, "azure-gpt-5.6-terra-messages_bridged_stream" +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_SOL_MESSAGES_BRIDGED_JSON, @@ -5903,6 +6015,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_LUNA_CACHE_WRITE_5M, AZURE_GPT_6_ASTRA_CACHE_WRITE_5M, AZURE_GPT_5_6_TERRA_INPUT_TEXT, + AZURE_GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K, + AZURE_GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_5_6_TERRA_STREAM, AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX, AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, @@ -5911,6 +6025,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M, AZURE_GPT_5_6_TERRA_REASONING, AZURE_GPT_5_6_TERRA_TOOL_CALL, + AZURE_GPT_6_1_SOL_RESPONSES_JSON, + AZURE_GPT_6_1_SOL_RESPONSES_STREAM, + AZURE_GPT_6_LUNA_RESPONSES_JSON, + AZURE_GPT_6_LUNA_RESPONSES_STREAM, + AZURE_GPT_6_ASTRA_RESPONSES_JSON, + AZURE_GPT_6_ASTRA_RESPONSES_STREAM, + AZURE_GPT_5_6_TERRA_RESPONSES_JSON, + AZURE_GPT_5_6_TERRA_RESPONSES_STREAM, + AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -5924,6 +6050,22 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_INPUT_TEXT, streamed=AZURE_GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_1_SOL_STREAM), StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_INPUT_TEXT, streamed=AZURE_GPT_5_6_TERRA_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_RESPONSES_JSON, streamed=AZURE_GPT_6_1_SOL_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_LUNA_RESPONSES_JSON, streamed=AZURE_GPT_6_LUNA_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_RESPONSES_JSON, streamed=AZURE_GPT_6_ASTRA_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_RESPONSES_JSON, streamed=AZURE_GPT_5_6_TERRA_RESPONSES_STREAM), + StreamParityTestCase( + plain=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -5938,4 +6080,10 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=AZURE_GPT_5_6_TERRA_INPUT_TEXT, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON ), + BridgeCostParityTestCase(native=AZURE_GPT_6_1_SOL_STREAM, bridged=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=AZURE_GPT_6_LUNA_STREAM, bridged=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=AZURE_GPT_6_ASTRA_STREAM, bridged=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase( + native=AZURE_GPT_5_6_TERRA_STREAM, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index f78bdfecfa3..858db06cafc 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -6088,6 +6088,71 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_WRITE_1H_ABOVE_100K: Final = with_usa ), ) +def _us_anthropic_messages_bridge_pairs( + input_case: CostTrackingTestCase, stream_case: CostTrackingTestCase +) -> tuple[BridgeCostParityTestCase, ...]: + return tuple( + BridgeCostParityTestCase( + native=pair.native, + bridged=pair.bridged.model_copy( + update={ + "name": ( + f"us.anthropic." + f"{pair.native.name.removeprefix('us.anthropic.').removesuffix('-input_text').removesuffix('-stream')}" + f"-messages_bridged_{'stream' if pair.native.request.get('stream') is True else 'json'}" + ) + } + ), + ) + for pair in bridge_pairs((input_case, stream_case), ("/v1/messages",)) + ) + + +def _us_anthropic_responses_stream_case( + json_case: CostTrackingTestCase, stream_case: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return json_case.model_copy( + update={ + "name": name, + "request": {**json_case.request, "stream": True}, + "response": stream_case.response, + } + ) + + +US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + "us.anthropic.claude-fable-5-1-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + "us.anthropic.claude-haiku-5-5-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + "us.anthropic.claude-opus-5-5-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + "us.anthropic.claude-sonnet-5-5-responses_bridged_stream", +) +_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM +) +_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM +) +_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM +) +_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_SONNET_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_SONNET_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, @@ -6210,6 +6275,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING, US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND, US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, + US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -6235,6 +6308,38 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS[1].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -6267,4 +6372,24 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + ), + *_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS, + *_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS, + *_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS, + *_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS, ) diff --git a/tests/integration/cost_calculation/cases/bedrock_mantle.py b/tests/integration/cost_calculation/cases/bedrock_mantle.py index 2fa940d66f4..3cd9150033a 100644 --- a/tests/integration/cost_calculation/cases/bedrock_mantle.py +++ b/tests/integration/cost_calculation/cases/bedrock_mantle.py @@ -18,7 +18,7 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -535,14 +535,40 @@ _MANTLE_GPT_CHAT_REQUEST: Final = { "messages": [{"role": "user", "content": "cost tracking fixture"}], "max_completion_tokens": 64, } +_MANTLE_GPT_TOOL_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ], + "tool_choice": "required", +} +_MANTLE_GPT_CACHE_WRITE_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "input": "cacheable pricing-harness prefix " * 1000, + "prompt_cache_key": "cost-cases-20261017", + "max_output_tokens": 64, +} def _mantle_gpt_usage( - input_tokens: int, output_tokens: int, reasoning_tokens: int, cached_tokens: int = 0 + input_tokens: int, + output_tokens: int, + reasoning_tokens: int, + cached_tokens: int = 0, + cache_write_tokens: int = 0, ) -> dict[str, JsonValue]: return { "input_tokens": input_tokens, - "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": cached_tokens}, + "input_tokens_details": {"cache_write_tokens": cache_write_tokens, "cached_tokens": cached_tokens}, "output_tokens": output_tokens, "output_tokens_details": {"reasoning_tokens": reasoning_tokens}, "total_tokens": input_tokens + output_tokens, @@ -628,6 +654,20 @@ def _mantle_gpt_cases( response_expected, _MANTLE_GPT_CHAT_REQUEST, ) + chat_stream_case: Final = chat_case.model_copy( + update={ + "name": f"bedrock_mantle-{provider_model}-chat_stream", + "request": {**chat_case.request, "stream": True}, + "response": stream_case.response, + } + ) + messages_pairs: Final = bridge_pairs((response_case, stream_case), ("/v1/messages",)) + messages_json: Final = messages_pairs[0].bridged.model_copy( + update={"name": f"bedrock_mantle-{provider_model}-messages_bridged_json"} + ) + messages_stream: Final = messages_pairs[1].bridged.model_copy( + update={"name": f"bedrock_mantle-{provider_model}-messages_bridged_stream"} + ) cache_case: Final = _mantle_gpt_json_case( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, provider_model, @@ -660,7 +700,73 @@ def _mantle_gpt_cases( tiered_boundary_expected, {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, ) - return (response_case, stream_case, reasoning_case, chat_case, cache_case, tiered_input, tiered_cache, tiered_boundary) + return ( + response_case, + stream_case, + reasoning_case, + chat_case, + cache_case, + tiered_input, + tiered_cache, + tiered_boundary, + chat_stream_case, + messages_json, + messages_stream, + ) + + +def _mantle_gpt_audit3_cases( + provider_model: str, + tool_input_tokens: int, + tool_output_tokens: int, + tool_reasoning_tokens: int, + tool_expected: ExactExpected, + cache_output_tokens: int, + cache_reasoning_tokens: int, + cache_expected: ExactExpected, +) -> tuple[CostTrackingTestCase, ...]: + tool_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tool_call", + _mantle_gpt_usage(tool_input_tokens, tool_output_tokens, tool_reasoning_tokens), + tool_expected, + _MANTLE_GPT_TOOL_REQUEST, + ) + if not isinstance(tool_case.response, JsonResponse): + raise TypeError(f"{tool_case.name}: expected a JSON response fixture") + tool_response: Final = tool_case.response.model_copy( + update={ + "body": { + **tool_case.response.body, + "output": [ + { + "arguments": '{"city":"Paris"}', + "call_id": "call_$REQUEST_ID", + "id": "fc_$REQUEST_ID", + "name": "get_weather", + "status": "completed", + "type": "function_call", + } + ], + } + } + ) + tool_case_with_response: Final = tool_case.model_copy(update={"response": tool_response}) + cache_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "cache_write_5m", + _mantle_gpt_usage( + 1667, + cache_output_tokens, + cache_reasoning_tokens, + cache_write_tokens=1665, + ), + cache_expected, + _MANTLE_GPT_CACHE_WRITE_REQUEST, + ) + return (tool_case_with_response, cache_case) def _mantle_gpt_bridge(case: CostTrackingTestCase) -> BridgeCostParityTestCase: @@ -723,6 +829,83 @@ _MANTLE_GPT_56_TERRA_CASES: Final = _mantle_gpt_cases( ExactExpected(spend=0.84537376, input_cost=0.83776, output_cost=0.00761376, prompt_tokens=272000, completion_tokens=412), ) +_MANTLE_GPT_AUDIT3_CASES: Final = ( + *_mantle_gpt_audit3_cases( + "openai.gpt-6.1-sol", + 50, + 20, + 0, + ExactExpected(spend=0.000363, input_cost=0.000121, output_cost=0.000242, prompt_tokens=50, completion_tokens=20), + 28, + 0, + ExactExpected( + spend=0.005381765, + input_cost=0.005042965, + output_cost=0.0003388, + prompt_tokens=1667, + completion_tokens=28, + cache_creation_cost=0.005038125, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-6-luna", + 51, + 18, + 0, + ExactExpected( + spend=0.000018612, input_cost=0.000006732, output_cost=0.00001188, prompt_tokens=51, completion_tokens=18 + ), + 64, + 40, + ExactExpected( + spend=0.000317229, + input_cost=0.000274989, + output_cost=0.00004224, + prompt_tokens=1667, + completion_tokens=64, + cache_creation_cost=0.000274725, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-6-astra", + 51, + 18, + 0, + ExactExpected( + spend=0.0020163, input_cost=0.0007293, output_cost=0.001287, prompt_tokens=51, completion_tokens=18 + ), + 27, + 0, + ExactExpected( + spend=0.031720975, + input_cost=0.029790475, + output_cost=0.0019305, + prompt_tokens=1667, + completion_tokens=27, + cache_creation_cost=0.029761875, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-5.6-terra", + 51, + 28, + 8, + ExactExpected( + spend=0.00067452, input_cost=0.00015708, output_cost=0.00051744, prompt_tokens=51, completion_tokens=28 + ), + 28, + 0, + ExactExpected( + spend=0.00693385, + input_cost=0.00641641, + output_cost=0.00051744, + prompt_tokens=1667, + completion_tokens=28, + cache_creation_cost=0.00641025, + ), + ), +) + _MANTLE_CLAUDE_OPUS_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" _ADAPTIVE_USAGE: Final = { "input_tokens": 24, @@ -737,7 +920,25 @@ _ADAPTIVE_EXPECTED: Final = ExactExpected( spend=0.0007524, input_cost=0.0001584, output_cost=0.000594, prompt_tokens=24, completion_tokens=18 ) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( + +def _mantle_without_usage_fields(case: CostTrackingTestCase) -> CostTrackingTestCase: + if not isinstance(case.response, JsonResponse): + raise TypeError(f"{case.name}: expected a JSON response fixture") + usage: Final = case.response.body.get("usage") + if not isinstance(usage, dict): + raise TypeError(f"{case.name}: expected a usage object") + response: Final = case.response.model_copy( + update={ + "body": { + **case.response.body, + "usage": {key: value for key, value in usage.items() if key not in {"inference_geo", "speed"}}, + } + } + ) + return case.model_copy(update={"response": response}) + + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_INPUT_TEXT, "bedrock_mantle-anthropic.claude-opus-5-5-input_text", _ADAPTIVE_USAGE, @@ -751,7 +952,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( "allowed_openai_params": ["reasoning_effort"], }, } -) +)) _MANTLE_CLAUDE_STREAM_RESPONSE: Final = SseResponse( content_type="text/event-stream", @@ -841,7 +1042,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES } ) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_CACHE_READ, "bedrock_mantle-anthropic.claude-opus-5-5-cache_read", { @@ -856,7 +1057,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( ExactExpected( spend=0.00308319, input_cost=0.00153219, output_cost=0.001551, prompt_tokens=4225, completion_tokens=47 ), -).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL})) BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = with_usage( CLAUDE_OPUS_5_5_CACHE_WRITE_5M, @@ -892,12 +1093,12 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = with_usage( ), ).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = with_usage( +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_THINKING, "bedrock_mantle-anthropic.claude-opus-5-5-thinking", _ADAPTIVE_USAGE, _ADAPTIVE_EXPECTED, -).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL})) BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL: Final = with_usage( CLAUDE_OPUS_5_5_TOOL_CALL, @@ -921,6 +1122,21 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON: Final = BEDROCK_MANTLE_CLAUDE_ } ) +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-responses_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-responses_stream", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line", "stream": True}, + } +) + BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( _mantle_gpt_bridge(_MANTLE_GPT_61_SOL_CASES[0]), _mantle_gpt_bridge(_MANTLE_GPT_6_LUNA_CASES[0]), @@ -930,6 +1146,14 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( native=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, bridged=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, ), + *tuple( + BridgeCostParityTestCase(native=cases[0], bridged=cases[9]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), + *tuple( + BridgeCostParityTestCase(native=cases[1], bridged=cases[10]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), ) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( @@ -945,11 +1169,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( *_MANTLE_GPT_6_LUNA_CASES, *_MANTLE_GPT_6_ASTRA_CASES, *_MANTLE_GPT_56_TERRA_CASES, + *_MANTLE_GPT_AUDIT3_CASES, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, @@ -972,4 +1199,16 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, ), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON, + streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM, + ), + *tuple( + StreamParityTestCase(plain=cases[3], streamed=cases[8]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), + *tuple( + StreamParityTestCase(plain=cases[9], streamed=cases[10]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), ) diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index bd64f4aeb11..ceca816c300 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -2,8 +2,10 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +import json from typing import Final, cast +from pydantic import JsonValue from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, @@ -17,26 +19,74 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: +def _gemini_case_response( + response: JsonResponse | SseResponse, + model_version: str, + usage_fields: dict[str, JsonValue] | None = None, +) -> JsonResponse | SseResponse: match response: case JsonResponse(): - return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + usage: Final = response.body.get("usageMetadata") + body: Final = {**response.body, "modelVersion": model_version} + return response.model_copy( + update={ + "body": { + **body, + **( + { + "usageMetadata": { + **cast(dict[str, JsonValue], usage), + **usage_fields, + } + } + if usage_fields is not None and isinstance(usage, dict) + else {} + ), + } + } + ) case SseResponse(): return response.model_copy( update={ - "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + "frames": tuple( + _gemini_case_frame(frame, model_version, usage_fields) for frame in response.frames + ) } ) +def _gemini_case_frame( + frame: str, model_version: str, usage_fields: dict[str, JsonValue] | None +) -> str: + replaced: Final = frame.replace("gemini-3.8-flash", model_version) + if usage_fields is None: + return replaced + prefix, separator, payload = replaced.partition("data: ") + if not separator or payload == "[DONE]": + return replaced + event: Final = cast(dict[str, JsonValue], json.loads(payload)) + usage: Final = event.get("usageMetadata") + if not isinstance(usage, dict): + return replaced + return f"{prefix}data: {json.dumps({**event, 'usageMetadata': {**usage, **usage_fields}})}" + + def _gemini_case( - sibling: CostTrackingTestCase, name: str, model: str, model_version: str, expected: ExactExpected + sibling: CostTrackingTestCase, + name: str, + model: str, + model_version: str, + expected: ExactExpected, + *, + usage_fields: dict[str, JsonValue] | None = None, ) -> CostTrackingTestCase: return sibling.model_copy( update={ "name": name, "model": model, - "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), model_version), + "response": _gemini_case_response( + cast(JsonResponse | SseResponse, sibling.response), model_version, usage_fields + ), "expected": expected, } ) @@ -1540,6 +1590,7 @@ GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1594,6 +1645,7 @@ GEMINI_GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1653,6 +1705,7 @@ GEMINI_GEMINI_3_8_FLASH_REASONING: Final = CostTrackingTestCase( "thoughtsTokenCount": 3480, "totalTokenCount": 5280, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1934,6 +1987,7 @@ GEMINI_GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -3275,19 +3329,46 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( GEMINI_GEMINI_4_ARGON_INPUT_TEXT: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-4-argon-input_text", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}]}, ) GEMINI_GEMINI_4_ARGON_STREAM: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-4-argon-stream", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}]}, ) GEMINI_GEMINI_4_ARGON_CACHE_READ: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-4-argon-cache_read", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380) + "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}]}, ) GEMINI_GEMINI_4_ARGON_REASONING: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-4-argon-reasoning", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040) + "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}]}, +) +GEMINI_GEMINI_4_ARGON_SERVICE_TIER_FLEX: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, + "gemini-gemini-4-argon-service_tier_flex", + "gemini/gemini-4-argon", + "gemini-4-argon", + ExactExpected(spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412), +) +GEMINI_GEMINI_4_ARGON_SERVICE_TIER_PRIORITY: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, + "gemini-gemini-4-argon-service_tier_priority", + "gemini/gemini-4-argon", + "gemini-4-argon", + ExactExpected(spend=0.015444, input_cost=0.0072864, output_cost=0.0081576, prompt_tokens=1840, completion_tokens=412), +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_SERVICE_TIER_PRIORITY: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, + "gemini-gemini-3.1-pro-preview-service_tier_priority", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=0.02018016, input_cost=0.0086112, output_cost=0.01156896, prompt_tokens=1840, completion_tokens=412 + ), ) GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-4-argon-tool_call", "gemini/gemini-4-argon", @@ -3296,28 +3377,60 @@ GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-3.1-pro-preview-input_text", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-3.1-pro-preview-stream", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-3.1-pro-preview-cache_read", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380) + ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-3.1-pro-preview-reasoning", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040) + ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-3.1-pro-preview-tool_call", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K: Final = _gemini_case( + GEMINI_GEMINI_3_1_PRO_TIERED_INPUT_ABOVE_200K, + "gemini-gemini-3.1-pro-preview-tiered_input_above_200k", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=1.079468, input_cost=1.06496, output_cost=0.014508, prompt_tokens=204800, completion_tokens=620 + ), + usage_fields={"serviceTier": "standard"}, +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K: Final = _gemini_case( + GEMINI_GEMINI_3_1_PRO_TIERED_CACHE_READ_ABOVE_200K, + "gemini-gemini-3.1-pro-preview-tiered_cache_read_above_200k", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=0.13742976, input_cost=0.12619776, output_cost=0.011232, prompt_tokens=205824, completion_tokens=480 + ), + usage_fields={"serviceTier": "standard"}, +) +GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_GEMINI_4_ARGON_INPUT_TEXT, GEMINI_GEMINI_4_ARGON_STREAM), + ("/v1/responses", "/v1/messages"), +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM), + ("/v1/responses", "/v1/messages"), +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_INPUT_TEXT, @@ -3379,12 +3492,19 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_4_ARGON_STREAM, GEMINI_GEMINI_4_ARGON_CACHE_READ, GEMINI_GEMINI_4_ARGON_REASONING, + GEMINI_GEMINI_4_ARGON_SERVICE_TIER_FLEX, + GEMINI_GEMINI_4_ARGON_SERVICE_TIER_PRIORITY, GEMINI_GEMINI_4_ARGON_TOOL_CALL, GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K, GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, + GEMINI_GEMINI_3_1_PRO_PREVIEW_SERVICE_TIER_PRIORITY, + *(pair.bridged for pair in GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES), + *(pair.bridged for pair in GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3410,6 +3530,22 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM ), + StreamParityTestCase( + plain=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[3].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[3].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3422,6 +3558,6 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( ), ("/v1/responses", "/v1/messages"), ), - *bridge_pairs((GEMINI_GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), - *bridge_pairs((GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), + *GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES, + *GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES, ) diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index f98dda5fb0c..760f77f91cc 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -19,13 +19,17 @@ from integration.cost_calculation.cost_tracking_case import ( RecountRates, SseResponse, WavUpload, + with_usage, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, @@ -33,12 +37,30 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_HAIKU_5_5_STREAM, CLAUDE_HAIKU_5_5_THINKING, CLAUDE_HAIKU_5_5_TOOL_CALL, ) +def _azure_ai_claude_messages_case( + response_case: CostTrackingTestCase, + request_case: CostTrackingTestCase, + name: str, + model: str, +) -> CostTrackingTestCase: + return response_case.model_copy( + update={ + "name": name, + "model": model, + "endpoint": "/v1/messages", + "request": request_case.request, + } + ) + + def _azure_ai_claude_usage( input_tokens: int, output_tokens: int, @@ -782,6 +804,31 @@ BEDROCK_INVOKE_HAIKU_STREAM: Final = CostTrackingTestCase( ), ) +BEDROCK_INVOKE_HAIKU_CHAT_STREAM: Final = CostTrackingTestCase( + name="bedrock-invoke-haiku-chat-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "bedrock-invoke-haiku-chat-stream"}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": 412, + }, + response=BEDROCK_INVOKE_HAIKU_STREAM.response, + expected=BEDROCK_INVOKE_HAIKU_STREAM.expected, +) + +BEDROCK_INVOKE_HAIKU_RESPONSES: Final = CostTrackingTestCase( + name="bedrock-invoke-haiku-responses", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "bedrock-invoke-haiku-responses", "max_output_tokens": 412}, + response=BEDROCK_INVOKE_HAIKU_JSON.response, + expected=BEDROCK_INVOKE_HAIKU_JSON.expected, +) + # azure_ai/gpt-5.4-mini-2026-03-17 AZURE_AI_GPT_5_4_MINI_LATEST: Final = CostTrackingTestCase( @@ -2002,7 +2049,66 @@ AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _azure_ai_claude_shape_case( ), ) - +AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + "azure_ai/claude-fable-5-1-messages_input_text", + "azure_ai/claude-fable-5-1", +) +AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, + "azure_ai/claude-fable-5-1-messages_stream", + "azure_ai/claude-fable-5-1", +) +AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES: Final = bridge_pairs( + (AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), + ("/v1/responses",), +) +AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + "azure_ai/claude-haiku-5-5-messages_input_text", + "azure_ai/claude-haiku-5-5", +) +AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + "azure_ai/claude-haiku-5-5-messages_stream", + "azure_ai/claude-haiku-5-5", +) +AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES: Final = bridge_pairs( + (AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( + AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + "azure_ai/claude-haiku-5-5-tiered_input_above_100k", + { + "input_tokens": 150000, + "output_tokens": 480, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + ExactExpected( + spend=0.14478, input_cost=0.1425, output_cost=0.00228, prompt_tokens=150000, completion_tokens=480 + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K: Final = with_usage( + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + "azure_ai/claude-haiku-5-5-tiered_cache_read_above_100k", + { + "input_tokens": 4096, + "output_tokens": 480, + "cache_read_input_tokens": 110592, + "cache_creation_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + ExactExpected( + spend=0.01667744, input_cost=0.01439744, output_cost=0.00228, prompt_tokens=114688, completion_tokens=480 + ), +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( PERPLEXITY_SONAR_NEXT_OFF_PEAK, NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND, @@ -2031,6 +2137,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( XAI_GROK_5_PROVIDER_REPORTED_COST, BEDROCK_INVOKE_HAIKU_JSON, BEDROCK_INVOKE_HAIKU_STREAM, + BEDROCK_INVOKE_HAIKU_CHAT_STREAM, + BEDROCK_INVOKE_HAIKU_RESPONSES, AZURE_AI_GPT_5_4_MINI_LATEST, AZURE_AI_GPT_5_4_MINI_LATEST_STREAM, AZURE_AI_GROK_4_6_REASONING_FOLDED_INTO_COMPLETION, @@ -2061,13 +2169,21 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, AZURE_AI_CLAUDE_FABLE_5_1_THINKING, AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + *(pair.bridged for pair in AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES), AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K, + AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, AZURE_AI_CLAUDE_HAIKU_5_5_THINKING, AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, + AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -2078,4 +2194,25 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=AZURE_AI_CLAUDE_HAIKU_5_5_STREAM ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + streamed=AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES[0].bridged, + streamed=AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES[1].bridged, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + streamed=AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES[0].bridged, + streamed=AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES[1].bridged, + ), +) + +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + *AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES, + *AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES, ) diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 85ac05dd69d..f610fd9591a 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -34,6 +34,256 @@ def _stream_for_model(response: SseResponse, model: str) -> SseResponse: return response.model_copy(update={"frames": tuple(frame.replace("gpt-6-luna", model) for frame in response.frames)}) +def _responses_json_body( + model: str, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> dict[str, JsonValue]: + reasoning_item: Final[dict[str, JsonValue]] = { + "content": [], + "encrypted_content": "$ENCRYPTED_CONTENT", + "id": "rs_$REQUEST_ID", + "summary": [], + "type": "reasoning", + } + message_item: Final[dict[str, JsonValue]] = { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + } + ], + "id": "msg_$REQUEST_ID", + "phase": "final_answer", + "role": "assistant", + "status": "completed", + "type": "message", + } + output: Final = [reasoning_item, message_item] if include_reasoning else [message_item] + body: Final[dict[str, JsonValue]] = { + "access_programs": [], + "background": False, + "completed_at": 1791855453, + "created_at": 1791855452, + "error": None, + "frequency_penalty": None, + "id": "resp_$REQUEST_ID", + "incomplete_details": None, + "instructions": None, + "max_output_tokens": None, + "max_tool_calls": None, + "metadata": {}, + "moderation": None, + "model": model, + "object": "response", + "output": output, + "parallel_tool_calls": True, + "presence_penalty": None, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": None, + "reasoning": {"effort": "medium"}, + "safety_identifier": None, + "service_tier": "default", + "status": "completed", + "store": True, + "temperature": 1.0, + "text": None, + "tool_choice": "auto", + "tool_usage": None, + "tools": [], + "top_logprobs": None, + "top_p": 1.0, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + } + return {**body, "content_filters": content_filters} if include_content_filters else body + + +def _responses_stream_response( + model: str, + delta_count: int, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> SseResponse: + response_body: Final = _responses_json_body( + model, + include_content_filters=include_content_filters, + content_filters=content_filters, + include_reasoning=include_reasoning, + ) + in_progress_body: Final = { + **response_body, + "output": [], + "status": "in_progress", + "usage": None, + } + reasoning_item: Final = { + "content": [], + "encrypted_content": "$ENCRYPTED_CONTENT", + "id": "rs_$REQUEST_ID", + "summary": [], + "type": "reasoning", + } + message_item: Final = { + "content": [], + "id": "msg_$REQUEST_ID", + "role": "assistant", + "status": "in_progress", + "type": "message", + } + delta_events: Final = tuple( + { + "content_index": 0, + "delta": "scripted response" if index == 0 else "", + "item_id": "msg_$REQUEST_ID", + "logprobs": [], + "obfuscation": None, + "output_index": 1 if include_reasoning else 0, + "sequence_number": (6 if include_reasoning else 4) + index, + "type": "response.output_text.delta", + } + for index in range(delta_count) + ) + delta_start: Final = 6 if include_reasoning else 4 + message_output_index: Final = 1 if include_reasoning else 0 + reasoning_events: Final = ( + ( + { + "item": reasoning_item, + "output_index": 0, + "sequence_number": 2, + "type": "response.output_item.added", + }, + { + "item": reasoning_item, + "output_index": 0, + "sequence_number": 3, + "type": "response.output_item.done", + }, + ) + if include_reasoning + else () + ) + completed_message_item: Final = { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + } + ], + "id": "msg_$REQUEST_ID", + "phase": "final_answer", + "role": "assistant", + "status": "completed", + "type": "message", + } + return SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "response": in_progress_body, + "sequence_number": 0, + "type": "response.created", + }, + { + "response": in_progress_body, + "sequence_number": 1, + "type": "response.in_progress", + }, + *reasoning_events, + { + "item": message_item, + "output_index": message_output_index, + "sequence_number": 4 if include_reasoning else 2, + "type": "response.output_item.added", + }, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "output_index": message_output_index, + "part": {"annotations": [], "logprobs": [], "text": "", "type": "output_text"}, + "sequence_number": 5 if include_reasoning else 3, + "type": "response.content_part.added", + }, + *delta_events, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "logprobs": [], + "output_index": message_output_index, + "sequence_number": delta_start + delta_count, + "text": "scripted response", + "type": "response.output_text.done", + }, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "output_index": message_output_index, + "part": { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + }, + "sequence_number": delta_start + delta_count + 1, + "type": "response.content_part.done", + }, + { + "item": completed_message_item, + "output_index": message_output_index, + "sequence_number": delta_start + delta_count + 2, + "type": "response.output_item.done", + }, + { + "response": response_body, + "sequence_number": delta_start + delta_count + 3, + "type": "response.completed", + }, + ), + ) + + +def _responses_stream_case( + base: CostTrackingTestCase, + name: str, + response_model: str, + delta_count: int, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> CostTrackingTestCase: + return base.model_copy( + update={ + "name": name, + "request": {**base.request, "stream": True}, + "response": _responses_stream_response( + response_model, + delta_count, + include_content_filters=include_content_filters, + content_filters=content_filters, + include_reasoning=include_reasoning, + ), + } + ) + + # gpt-5.3-codex GPT_5_3_CODEX_INPUT_TEXT: Final = CostTrackingTestCase( name="gpt-5.3-codex-input_text", @@ -11356,6 +11606,31 @@ GPT_5_6_TERRA_TOOL_CALL: Final = GPT_6_LUNA_TOOL_CALL.model_copy( } ) +GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( + GPT_5_6_TERRA_INPUT_TEXT, + "gpt-5.6-terra-tiered_input_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + }, + ExactExpected( + spend=1.5696408, input_cost=1.56, output_cost=0.0096408, prompt_tokens=300000, completion_tokens=412 + ), +) +GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K: Final = with_usage( + GPT_5_6_TERRA_INPUT_TEXT, + "gpt-5.6-terra-tiered_cache_read_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, + }, + ExactExpected( + spend=0.6336408, input_cost=0.624, output_cost=0.0096408, prompt_tokens=300000, completion_tokens=412 + ), +) + GPT_5_6_SLASH_ALIAS: Final = GPT_5_6_INPUT_TEXT.model_copy( update={"name": "gpt-5.6-slash_alias", "alias": "openai/custom/gpt-5.6-$MARKER"} @@ -11446,6 +11721,63 @@ GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS: Final = CostTrackingTestCase( ), ) +GPT_6_1_SOL_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_1_SOL_RESPONSES_NATIVE_JSON, + "gpt-6.1-sol-responses_native_stream", + "gpt-6.1-sol", + 1, + include_reasoning=False, +) +GPT_6_LUNA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_LUNA_RESPONSES_NATIVE_JSON, + "gpt-6-luna-responses_native_stream", + "gpt-6-luna", + 45, +) +GPT_6_ASTRA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_ASTRA_RESPONSES_NATIVE_JSON, + "gpt-6-astra-responses_native_stream", + "gpt-6-astra", + 36, +) +GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, + "gpt-5.6-terra-responses_native_stream", + "gpt-5.6-terra", + 95, +) +GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON: Final = bridge_pairs( + (GPT_5_6_TERRA_RESPONSES_NATIVE_JSON,), ("/v1/messages",) +)[0].bridged.model_copy(update={"name": "gpt-5.6-terra-messages_bridged_json"}) +GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM: Final = GPT_6_1_SOL_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6.1-sol-messages_bridged_stream", + "request": {**GPT_6_1_SOL_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_1_SOL_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_6_LUNA_MESSAGES_BRIDGED_STREAM: Final = GPT_6_LUNA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6-luna-messages_bridged_stream", + "request": {**GPT_6_LUNA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_LUNA_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM: Final = GPT_6_ASTRA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6-astra-messages_bridged_stream", + "request": {**GPT_6_ASTRA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_ASTRA_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-messages_bridged_stream", + "request": {**GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM.response, + } +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, @@ -11659,6 +11991,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_6_LUNA_CACHE_WRITE_5M, GPT_6_ASTRA_CACHE_WRITE_5M, GPT_5_6_TERRA_INPUT_TEXT, + GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K, + GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K, GPT_5_6_TERRA_STREAM, GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, @@ -11669,6 +12003,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_TERRA_SERVICE_TIER_FLEX, GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, GPT_5_6_TERRA_TOOL_CALL, + GPT_6_1_SOL_RESPONSES_NATIVE_STREAM, + GPT_6_LUNA_RESPONSES_NATIVE_STREAM, + GPT_6_ASTRA_RESPONSES_NATIVE_STREAM, + GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM, + GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, + GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM, + GPT_6_LUNA_MESSAGES_BRIDGED_STREAM, + GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM, + GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -11698,6 +12041,16 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=GPT_6_ASTRA_INPUT_TEXT, streamed=GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=GPT_6_1_SOL_INPUT_TEXT, streamed=GPT_6_1_SOL_STREAM), StreamParityTestCase(plain=GPT_5_6_TERRA_INPUT_TEXT, streamed=GPT_5_6_TERRA_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, streamed=GPT_6_1_SOL_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_RESPONSES_NATIVE_JSON, streamed=GPT_6_LUNA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, streamed=GPT_6_ASTRA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, streamed=GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, streamed=GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_MESSAGES_BRIDGED_JSON, streamed=GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, streamed=GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase( + plain=GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, streamed=GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -11726,6 +12079,13 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_JSON, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, bridged=GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_STREAM, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_STREAM, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_STREAM, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase( + native=GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM, bridged=GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), BridgeCostParityTestCase(native=GPT_6_SOL_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_SOL_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_LUNA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_ASTRA_CHAT_BRIDGED_JSON), diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 761078ef593..0e95eb865bf 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -2,6 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +import json from typing import Final, cast from pydantic import JsonValue @@ -11,13 +12,19 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, CLAUDE_OPUS_5_5_CACHE_WRITE_1H, CLAUDE_OPUS_5_5_CACHE_WRITE_5M, CLAUDE_OPUS_5_5_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_STREAM, CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, @@ -25,6 +32,8 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_SONNET_5_5_CACHE_WRITE_1H, CLAUDE_SONNET_5_5_CACHE_WRITE_5M, CLAUDE_SONNET_5_5_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_STREAM, CLAUDE_SONNET_5_5_STREAM, CLAUDE_SONNET_5_5_THINKING, CLAUDE_SONNET_5_5_TOOL_CALL, @@ -41,26 +50,72 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: +def _gemini_case_response( + response: JsonResponse | SseResponse, + model_version: str, + usage_fields: dict[str, JsonValue] | None = None, +) -> JsonResponse | SseResponse: match response: case JsonResponse(): - return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + usage: Final = response.body.get("usageMetadata") + body: Final = {**response.body, "modelVersion": model_version} + return response.model_copy( + update={ + "body": { + **body, + **( + { + "usageMetadata": { + **cast(dict[str, JsonValue], usage), + **usage_fields, + } + } + if usage_fields is not None and isinstance(usage, dict) + else {} + ), + } + } + ) case SseResponse(): return response.model_copy( update={ - "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + "frames": tuple( + _gemini_case_frame(frame, model_version, usage_fields) for frame in response.frames + ) } ) +def _gemini_case_frame(frame: str, model_version: str, usage_fields: dict[str, JsonValue] | None) -> str: + replaced: Final = frame.replace("gemini-3.8-flash", model_version) + if usage_fields is None: + return replaced + prefix, separator, payload = replaced.partition("data: ") + if not separator or payload == "[DONE]": + return replaced + event: Final = cast(dict[str, JsonValue], json.loads(payload)) + usage: Final = event.get("usageMetadata") + if not isinstance(usage, dict): + return replaced + return f"{prefix}data: {json.dumps({**event, 'usageMetadata': {**usage, **usage_fields}})}" + + def _gemini_case( - sibling: CostTrackingTestCase, name: str, model: str, expected: ExactExpected + sibling: CostTrackingTestCase, + name: str, + model: str, + expected: ExactExpected, + *, + usage_fields: dict[str, JsonValue] | None = None, + model_version: str = "gemini-3.1-pro-preview", ) -> CostTrackingTestCase: return sibling.model_copy( update={ "name": name, "model": model, - "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), "gemini-3.1-pro-preview"), + "response": _gemini_case_response( + cast(JsonResponse | SseResponse, sibling.response), model_version, usage_fields + ), "expected": expected, } ) @@ -97,6 +152,7 @@ def _vertex_claude_shape_case( expected: ExactExpected, response_content: list[dict[str, JsonValue]] | None = None, adaptive_thinking: bool = False, + thinking_effort: str = "medium", ) -> CostTrackingTestCase: response: Final = cast(JsonResponse, sibling.response) request: Final = ( @@ -107,7 +163,7 @@ def _vertex_claude_shape_case( if key not in {"reasoning_effort", "allowed_openai_params"} }, "thinking": {"type": "adaptive"}, - "output_config": {"effort": "medium"}, + "output_config": {"effort": thinking_effort}, "allowed_openai_params": ["thinking", "output_config"], } if adaptive_thinking @@ -1716,6 +1772,8 @@ GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -1770,6 +1828,8 @@ GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -2117,6 +2177,8 @@ GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -3122,28 +3184,258 @@ VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( ), ) +VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT, + "vertex_ai/claude-haiku-5-5-input_text", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=1840, output_tokens=412), + ExactExpected(spend=0.00507, input_cost=0.002392, output_cost=0.002678, prompt_tokens=1840, completion_tokens=412), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM: Final = VERTEX_AI_CLAUDE_SONNET_5_5_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-stream", + "model": "vertex_ai/claude-haiku-5-5", + "response": _vertex_claude_stream_response("claude-haiku-5-5", "end_turn", ("scripted answer",)), + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, + "vertex_ai/claude-haiku-5-5-cache_read", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.00356759, + input_cost=0.00352859, + output_cost=0.000039, + prompt_tokens=27017, + completion_tokens=6, + cache_read_cost=0.00351039, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-haiku-5-5-cache_write_5m", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.043950075, + input_cost=0.043911075, + output_cost=0.000039, + prompt_tokens=27025, + completion_tokens=6, + cache_creation_cost=0.043892875, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-haiku-5-5-cache_write_1h", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.0702858, + input_cost=0.0702468, + output_cost=0.000039, + prompt_tokens=27025, + completion_tokens=6, + cache_creation_cost=0.0702286, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_THINKING: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_THINKING, + "vertex_ai/claude-haiku-5-5-thinking", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=144, thinking_tokens=37), + ExactExpected( + spend=0.0009594, input_cost=0.0000234, output_cost=0.000936, prompt_tokens=18, completion_tokens=144 + ), + adaptive_thinking=True, + thinking_effort="low", +) +VERTEX_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL, + "vertex_ai/claude-haiku-5-5-tool_call", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.0008255, input_cost=0.0005005, output_cost=0.000325, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-messages_input_text", + "model": "vertex_ai/claude-haiku-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = CLAUDE_HAIKU_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-messages_stream", + "model": "vertex_ai/claude-haiku-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) + +VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-1-messages_input_text", + "model": "vertex_ai/claude-fable-5-1", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-1-messages_stream", + "model": "vertex_ai/claude-fable-5-1", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_1_STREAM.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-opus-5-5-messages_input_text", + "model": "vertex_ai/claude-opus-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-opus-5-5-messages_stream", + "model": "vertex_ai/claude-opus-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_OPUS_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-sonnet-5-5-messages_input_text", + "model": "vertex_ai/claude-sonnet-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM: Final = CLAUDE_SONNET_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-sonnet-5-5-messages_stream", + "model": "vertex_ai/claude-sonnet-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_SONNET_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) + GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_3_8_FLASH_INPUT_TEXT, "vertex_ai/gemini-3.1-pro-preview-input_text", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( GEMINI_3_8_FLASH_STREAM, "vertex_ai/gemini-3.1-pro-preview-stream", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( GEMINI_3_8_FLASH_CACHE_READ, "vertex_ai/gemini-3.1-pro-preview-cache_read", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) + ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( GEMINI_3_1_PRO_REASONING, "vertex_ai/gemini-3.1-pro-preview-reasoning", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) + ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}], + "trafficType": "ON_DEMAND", + }, +) +GEMINI_3_8_FLASH_REASONING: Final = _gemini_case( + GEMINI_3_1_PRO_REASONING, + "vertex_ai/gemini-3.8-flash-reasoning", + "vertex_ai/gemini-3.8-flash", + ExactExpected( + spend=0.01144, + input_cost=0.0006448, + output_cost=0.0107952, + prompt_tokens=1240, + completion_tokens=4040, + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}], + "trafficType": "ON_DEMAND", + }, + model_version="gemini-3.8-flash", ) GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( GEMINI_3_8_FLASH_TOOL_CALL, "vertex_ai/gemini-3.1-pro-preview-tool_call", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K: Final = _gemini_case( + GEMINI_3_1_PRO_TIERED_INPUT_ABOVE_200K, + "vertex_ai/gemini-3.1-pro-preview-tiered_input_above_200k", + "vertex_ai/gemini-3.1-pro-preview", + ExactExpected( + spend=0.913396, input_cost=0.90112, output_cost=0.012276, prompt_tokens=204800, completion_tokens=620 + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 620}], + "trafficType": "ON_DEMAND", + }, +) +GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K: Final = _gemini_case( + GEMINI_3_1_PRO_TIERED_CACHE_READ_ABOVE_200K, + "vertex_ai/gemini-3.1-pro-preview-tiered_cache_read_above_200k", + "vertex_ai/gemini-3.1-pro-preview", + ExactExpected( + spend=0.11628672, input_cost=0.10678272, output_cost=0.009504, prompt_tokens=205824, completion_tokens=480 + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 480}], + "trafficType": "ON_DEMAND", + "cacheTokensDetails": [{"modality": "TEXT", "tokenCount": 201728}], + }, +) + GEMINI_4_ARGON_INPUT_TEXT: Final = GEMINI_3_8_FLASH_INPUT_TEXT.model_copy( update={ "name": "vertex_ai/gemini-4-argon-input_text", @@ -3301,6 +3593,15 @@ GEMINI_4_ARGON_TOOL_CALL: Final = GEMINI_3_8_FLASH_TOOL_CALL.model_copy( } ) +GEMINI_4_ARGON_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_4_ARGON_INPUT_TEXT, GEMINI_4_ARGON_STREAM), + ("/v1/responses", "/v1/messages"), +) +GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM), + ("/v1/responses", "/v1/messages"), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_INPUT_TEXT, GEMINI_3_1_PRO_CACHE_READ, @@ -3326,6 +3627,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_STREAM_TOOL_CALL, GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_3_8_FLASH_INPUT_TEXT, + GEMINI_3_8_FLASH_REASONING, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, @@ -3367,9 +3669,30 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_HAIKU_5_5_THINKING, + VERTEX_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS), GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K, + GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K, GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, GEMINI_4_ARGON_INPUT_TEXT, @@ -3377,6 +3700,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_4_ARGON_CACHE_READ, GEMINI_4_ARGON_REASONING, GEMINI_4_ARGON_TOOL_CALL, + *(pair.bridged for pair in GEMINI_4_ARGON_ROUTE_BRIDGES), + *(pair.bridged for pair in GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3397,8 +3722,59 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS[1].bridged, + ), StreamParityTestCase(plain=GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_3_1_PRO_PREVIEW_STREAM), StreamParityTestCase(plain=GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_4_ARGON_STREAM), + StreamParityTestCase( + plain=GEMINI_4_ARGON_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_4_ARGON_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_4_ARGON_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_4_ARGON_ROUTE_BRIDGES[3].bridged, + ), + StreamParityTestCase( + plain=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[3].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3406,6 +3782,10 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), ("/v1/responses", "/v1/messages"), ), - *bridge_pairs((GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), - *bridge_pairs((GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), + *GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES, + *GEMINI_4_ARGON_ROUTE_BRIDGES, + *VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS, ) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 2ef42249a7e..ba7be7647da 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -889,6 +889,26 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "max_output_tokens": 128000, "max_tokens": 128000, }, + "vertex_ai/claude-haiku-5-5": { + "litellm_provider": "vertex_ai-anthropic_models", + "mode": "chat", + "input_cost_per_token": 1.3e-06, + "input_cost_per_token_above_100k_tokens": 6.5e-06, + "output_cost_per_token": 6.5e-06, + "output_cost_per_token_above_100k_tokens": 3.25e-05, + "cache_creation_input_token_cost": 1.625e-06, + "cache_creation_input_token_cost_above_1hr": 2.6e-06, + "cache_creation_input_token_cost_above_100k_tokens": 8.125e-06, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.3e-05, + "cache_read_input_token_cost": 1.3e-07, + "cache_read_input_token_cost_above_100k_tokens": 6.5e-07, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + }, "groq/qwen/qwen3.8-27b": { "litellm_provider": "groq", "mode": "chat", diff --git a/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py new file mode 100644 index 00000000000..8af70d39186 --- /dev/null +++ b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py @@ -0,0 +1,168 @@ +"""A native Responses WebSocket runs the proxy's pre_call guardrails on every ``response.create`` frame, bills each +turn's spend row with the guardrail cost of that turn alone, and bills a blocked frame's guardrail scan as a failed +turn (LIT-9744)""" + +import asyncio +import hashlib +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import websockets +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, create, responses_peer, turn +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_INPUT_COST_PER_TOKEN: Final = 0.002 +_OUTPUT_COST_PER_TOKEN: Final = 0.004 +_COST_PER_WORD: Final = 0.1 +_MARKER: Final = " [guarded]" +_BLOCKED_WORD: Final = "forbidden" +_GUARDRAIL_MODULE: Final = f""" +from litellm.integrations.custom_guardrail import CustomGuardrail + + +class TurnGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + blocks = [block for item in data.get("input") or [] for block in item.get("content") or []] + words = sum(len(block["text"].split()) for block in blocks) + blocked = any({_BLOCKED_WORD!r} in block["text"] for block in blocks) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={{"words": words}}, + request_data=data, + guardrail_status="guardrail_intervened" if blocked else "success", + tracing_detail={{"guardrail_cost": round({_COST_PER_WORD} * words, 10)}}, + ) + if blocked: + raise ValueError("turn guardrail blocked the request") + for block in blocks: + block["text"] = block["text"] + {_MARKER!r} + return data +""" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + cert: tuple[Path, Path] + + +def _config(directory: Path) -> Path: + (directory / "turn_guardrail.py").write_text(_GUARDRAIL_MODULE) + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["guardrails"] = [ + { + "guardrail_name": "turn-guardrail", + "litellm_params": {"guardrail": "turn_guardrail.TurnGuardrail", "mode": "pre_call", "default_on": True}, + } + ] + path: Final = directory / "responses-ws-guardrail.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("responses-ws-guardrail") + cert: Final = write_self_signed_cert(directory) + with ( + gateway_from_environment() as shared, + owned_proxy(shared, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_config(directory)) as owned, + ): + yield _Rig(owned, cert) + + +def _nested(frame: str) -> str: + """The ``{"type": "response.create", "response": {...}}`` wire shape the WS layer accepts next to the flat one.""" + flat: Final = json.loads(frame) + return json.dumps({"type": flat["type"], "response": {key: value for key, value in flat.items() if key != "type"}}) + + +async def _session( + proxy_url: str, key: str, model: str, frames: tuple[str, ...] +) -> tuple[tuple[dict[str, JsonValue], ...], ...]: + proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/responses?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=10, + ) as connection: + return tuple([await turn(connection, frame) for frame in frames]) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT spend, status FROM "LiteLLM_SpendLogs" WHERE api_key = %s ORDER BY "startTime"', + (hashlib.sha256(key.encode()).hexdigest(),), + ) + + +def _key_spend(key: str) -> float: + (row,) = read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashlib.sha256(key.encode()).hexdigest(),) + ) + return float(row["spend"]) + + +def _llm_cost() -> float: + return USAGE["input_tokens"] * _INPUT_COST_PER_TOKEN + USAGE["output_tokens"] * _OUTPUT_COST_PER_TOKEN + + +def test_every_turn_runs_the_pre_call_guardrail_and_bills_only_its_own_guardrail_cost(rig: _Rig) -> None: + with responses_peer(rig.cert) as peer, rig.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=_INPUT_COST_PER_TOKEN, + output_cost_per_token=_OUTPUT_COST_PER_TOKEN, + ) + key: Final = scenario.key(models=[model]) + tag: Final = uuid.uuid4().hex + texts: Final = ( + f"one {tag}", + f"one two {tag}", + f"one two three {tag}", + f"{_BLOCKED_WORD} {tag}", + f"one {tag}", + f"{_BLOCKED_WORD} again {tag}", + ) + frames: Final = tuple(create(model, text) for text in texts[:-1]) + (_nested(create(model, texts[-1])),) + turns: Final = asyncio.run(_session(str(rig.gateway.client.base_url), key, model, frames)) + + assert [events[-1]["type"] for events in turns] == [ + "response.completed", + "response.completed", + "response.completed", + "error", + "response.completed", + "error", + ], [[event["type"] for event in events] for events in turns] + assert "blocked" in str(turns[3][-1]["error"]), turns[3] + + forwarded: Final = tuple(peer.frames.get_nowait() for _ in range(peer.frames.qsize())) + assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == [ + f"{text}{_MARKER}" for text in texts if _BLOCKED_WORD not in text + ], forwarded + + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 6, seconds=70) + assert [row["status"] for row in rows] == ["success", "success", "success", "failure", "success", "failure"], ( + rows + ) + expected: Final = [ + (0.0 if _BLOCKED_WORD in text else _llm_cost()) + _COST_PER_WORD * len(text.split()) for text in texts + ] + assert [float(row["spend"]) for row in rows] == pytest.approx(expected), rows + total: Final = sum(expected) + assert eventually( + lambda: _key_spend(key), lambda spend: abs(spend - total) < 1e-6, seconds=70 + ) == pytest.approx(total) diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index 1c0cf42ad13..b7bf3cfe968 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -661,10 +661,6 @@ def test_scim_admin_group_put_members_promotes_and_demotes(gateway: Gateway, tmp @pytest.mark.timeout(300) def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip( - "BUG: PUT /scim/v2/Groups/{id} returns the members list from before the update, so the 200 body omits " - "the members it just added" - ) group_name: Final = f"scim-admin-put-response-{uuid.uuid4().hex}" with _scim_admin_group_fixture(gateway, tmp_path, group_name) as fixture: expected_member: Final = ScimGroupMember( diff --git a/tests/integration/management/test_scim_group_bulk_roster.py b/tests/integration/management/test_scim_group_bulk_roster.py index 8387beca958..2af6ae0a83d 100644 --- a/tests/integration/management/test_scim_group_bulk_roster.py +++ b/tests/integration/management/test_scim_group_bulk_roster.py @@ -18,7 +18,6 @@ from pathlib import Path from typing import Final, LiteralString import httpx -import psutil import psycopg import pytest from integration._support.budget import case_timeout_seconds @@ -55,9 +54,6 @@ LOCK_KEY_ROW_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE t SESSIONS_BLOCKED_BY_SQL: Final = ( "SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))" ) -CLIENT_PORTS_BLOCKED_BY_SQL: Final = ( - "SELECT client_port FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid)) AND client_port > 0" -) TRANSACTIONS_OPEN_SINCE_SQL: Final = """ SELECT count(*)::int AS open FROM pg_stat_activity WHERE datname = current_database() AND backend_type = 'client backend' AND pid <> pg_backend_pid() @@ -296,22 +292,6 @@ def _transactions_open_since(moment: str) -> int: return _int_value(read_rows(TRANSACTIONS_OPEN_SINCE_SQL, (moment,)), "open") -def _client_ports_blocked_by(backend_pid: int) -> frozenset[int]: - rows: Final = read_rows(CLIENT_PORTS_BLOCKED_BY_SQL, (str(backend_pid),)) - return frozenset(_int_value((row,), "client_port") for row in rows) - - -def _local_tcp_ports(pid: int) -> Iterator[int]: - worker: Final = psutil.Process(pid) - for process in (worker, *worker.children(recursive=True)): - yield from (connection.laddr.port for connection in process.net_connections(kind="tcp")) - - -def _worker_holding(owned: OwnedProxy, client_ports: frozenset[int]) -> int: - (worker,) = tuple(pid for pid in owned.worker_pids() if client_ports & frozenset(_local_tcp_ports(pid))) - return worker - - def _roster_ids(team: str) -> tuple[str, ...]: rows: Final = read_rows(TEAM_ROW_SQL, (team,)) assert len(rows) == 1, rows @@ -508,11 +488,16 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne pinned_key: Final = _team_key_for(scenario, team, max(leaving)) with _holding_row_lock(LOCK_KEY_ROW_SQL, sha256(pinned_key.encode()).hexdigest()) as holder_pid: push: Final = pool.submit(_push_outcome, candidate, team, arriving) - blocked_ports: Final = eventually( - lambda: _client_ports_blocked_by(holder_pid), bool, seconds=UNFIXED_PUSH_SECONDS + _ = eventually( + lambda: _sessions_blocked_by(holder_pid), + lambda waiting: waiting > 0, + seconds=UNFIXED_PUSH_SECONDS, ) killed_at: Final = _database_now() - _ = owned.kill_worker(_worker_holding(owned, blocked_ports)) + victims: Final = tuple(sorted(owned.worker_pids())) + assert victims, "expected at least one serving worker" + for pid in victims: + _ = owned.kill_worker(pid) record_property("interrupted_push", push.result()) _ = eventually(lambda: _transactions_open_since(killed_at), lambda open: open == 0, seconds=ROLLBACK_SECONDS) settled: Final = frozenset(_roster_ids(team)) diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 9d15ea0afa9..6f59f694352 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -860,6 +860,7 @@ def test_spec_client_discovers_the_aggregate_authorization_server_and_signs_in_t "client_origin": "http://127.0.0.1:9", "server_id": None, "server_name": None, + "servers": None, "connected": None, }, described.text completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies) @@ -969,6 +970,7 @@ def test_per_server_resource_scopes_the_session_bearer_to_the_named_server(gatew "client_origin": "http://127.0.0.1:9", "server_id": server_a, "server_name": alias_a, + "servers": None, "connected": True, }, described.text completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py new file mode 100644 index 00000000000..e2f23b7558f --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py @@ -0,0 +1,348 @@ +import asyncio +import json +import re +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.messages_endpoint.providers.anthropic.test_anthropic_safeguards_beta_wire import ( + ANTHROPIC_KEY, + ANTHROPIC_MODEL, + BETA, + OMITTED, + anthropic_peer, + marker_of, + messages_body, + stream_frames, +) +from pydantic import JsonValue, TypeAdapter + +_CONFIG_MODEL: Final = "anthropic-safeguards-chaos" +_FOUNDRY_CONFIG_MODEL: Final = "azure-ai-safeguards-chaos" +_YAML_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MODEL_LIST: Final = json.dumps( + {"object": "list", "data": [{"id": ANTHROPIC_MODEL, "object": "model", "owned_by": "anthropic"}]} +).encode() + + +@dataclass(frozen=True, slots=True) +class _Call: + stream: bool + marker: str + safeguards: bool = True + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _calls(count: int, stream: bool, safeguards: bool = True) -> tuple[_Call, ...]: + return tuple(_Call(stream=stream, marker=uuid.uuid4().hex, safeguards=safeguards) for _ in range(count)) + + +def _streamed_reply(marker: str, abort_after: int | None = None, pause: float = 0) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=stream_frames(marker), + abort_after=abort_after, + pause_between_chunks=pause, + ) + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + "/v1/messages", + json=( + messages_body(model, call.marker, stream=call.stream) + if call.safeguards + else messages_body(model, call.marker, stream=call.stream, safeguards=OMITTED) + ), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + assert f"msg_{served.call.marker}" in served.text, served.text + if served.call.stream: + assert "event: message_stop" in served.text, served.text + + +def _assert_every_request_carried_the_beta(received: tuple[Request, ...], markers: frozenset[str]) -> None: + assert sorted(marker_of(request) for request in received) == sorted(markers) + missing: Final = [ + frozenset(request.headers) for request in received if request.headers.get("anthropic-beta") != BETA + ] + assert missing == [], missing + + +def _spend_success_ids(model: str, expected: frozenset[str]) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s', (model, "success") + ), + lambda found: expected <= {str(row["request_id"]) for row in found}, + seconds=70, + ) + identities: Final = [row["request_id"] for row in rows] + assert len(set(identities)) == len(identities), identities + + +async def test_burst_with_aborted_streams_keeps_the_beta_on_every_forwarded_request(gateway: Gateway) -> None: + calls: Final = tuple(_Call(stream=index % 2 == 0, marker=uuid.uuid4().hex) for index in range(30)) + aborted: Final = frozenset(call.marker for index, call in enumerate(calls) if call.stream and index % 3 == 0) + + def respond(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker in aborted: + return _streamed_reply(marker, abort_after=0) + return anthropic_peer(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 30 + for item in served: + if item.call.marker in aborted: + assert item.status == 500, item.text + assert '"type":"api_error"' in item.text and "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(stream=True, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in (*calls, recovery))) + answered: Final = frozenset(f"msg_{call.marker}" for call in (*calls, recovery) if call.marker not in aborted) + _spend_success_ids(model, answered) + + +async def test_slow_streams_complete_with_the_beta_on_every_request(gateway: Gateway) -> None: + calls: Final = _calls(10, stream=True) + with ( + wire_server(lambda request: _streamed_reply(marker_of(request), pause=0.3)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 10 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in calls)) + _spend_success_ids(model, frozenset(f"msg_{call.marker}" for call in calls)) + + +def _anthropic_deployment(wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": ANTHROPIC_KEY}, + } + + +def _foundry_deployment(wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": _FOUNDRY_CONFIG_MODEL, + "litellm_params": { + "model": "azure_ai/claude-sonnet-4-6", + "api_base": wire.url, + "api_key": "synthetic-foundry-key", + }, + } + + +def _chaos_config(tmp_path: Path, *deployments: dict[str, JsonValue]) -> Path: + base: Final = _YAML_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = {**base, "model_list": list(deployments)} + path: Final = tmp_path / "anthropic-safeguards-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(180) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_sending_the_beta(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(20, stream=False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if (request.method, request.target) == ("GET", "/v1/models"): + return Reply(body=_MODEL_LIST) + held_markers.put(marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return anthropic_peer(request) + + with wire_server(held) as wire: + path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + received: Final = wire.drain() + posts: Final = tuple(request for request in received if request.method == "POST") + probes: Final[list[tuple[str, str]]] = [ + (request.method, request.target) for request in received if request.method != "POST" + ] + assert set(probes) <= {("GET", "/v1/models")}, probes + _assert_every_request_carried_the_beta(posts, frozenset(call.marker for call in (*calls, follow_up))) + + +def _beta_map_gets(received: tuple[Request, ...]) -> int: + return sum(1 for request in received if (request.method, request.target) == ("GET", "/beta.json")) + + +async def _chat(base_url: str, key: str, marker: str) -> httpx.Response: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + return await client.post( + "/v1/chat/completions", + json={ + "model": _CONFIG_MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Question marker-{marker}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + + +async def _answered_while_the_beta_map_is_held( + base_url: str, key: str, guarded: tuple[_Call, ...], chat_markers: tuple[str, ...] +) -> None: + served: Final = await _burst(base_url, key, _CONFIG_MODEL, guarded) + chats: Final = await asyncio.gather(*(_chat(base_url, key, marker) for marker in chat_markers)) + for item in served: + _assert_answered_with_its_own_marker(item) + for marker, chat in zip(chat_markers, chats, strict=True): + assert chat.status_code == 200, chat.text + assert f"marker-{marker}" in chat.text, chat.text + + +@pytest.mark.timeout(240) +async def test_cold_workers_fetch_the_beta_map_only_for_filtered_safeguards_requests_off_the_shared_threads( + gateway: Gateway, tmp_path: Path +) -> None: + beta_map: Final = Path("litellm/anthropic_beta_headers_config.json").read_bytes() + release: Final = threading.Event() + + def held_beta_map(_: Request) -> Reply: + assert release.wait(timeout=180), "The beta map was never released" + return Reply(body=beta_map) + + def upstream(request: Request) -> Reply: + if (request.method, request.target) == ("GET", "/v1/models"): + return Reply(body=_MODEL_LIST) + return anthropic_peer(request) + + plain: Final = _calls(8, stream=False, safeguards=False) + plain_foundry: Final = _calls(8, stream=False, safeguards=False) + direct: Final = _calls(8, stream=False) + foundry: Final = _calls(80, stream=False) + direct_while_held: Final = _calls(8, stream=False) + chat_markers: Final = tuple(uuid.uuid4().hex for _ in range(4)) + with wire_server(held_beta_map) as beta_map_wire, wire_server(upstream) as wire: + path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire), _foundry_deployment(wire)) + with owned_proxy_process( + gateway, + tmp_path, + {"LITELLM_ANTHROPIC_BETA_HEADERS_URL": f"{beta_map_wire.url}/beta.json"}, + remove_environment=("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS",), + config=path, + workers=2, + ) as owned: + base_url: Final = str(owned.gateway.client.base_url) + key: Final = owned.gateway.key + cold: Final = ( + *await _burst(base_url, key, _CONFIG_MODEL, plain), + *await _burst(base_url, key, _FOUNDRY_CONFIG_MODEL, plain_foundry), + *await _burst(base_url, key, _CONFIG_MODEL, direct), + ) + for served_cold in cold: + _assert_answered_with_its_own_marker(served_cold) + assert beta_map_wire.received.qsize() == 0, "a request that never filters betas fetched the beta map" + + held: Final = asyncio.create_task(_burst(base_url, key, _FOUNDRY_CONFIG_MODEL, foundry)) + try: + await asyncio.to_thread(eventually, beta_map_wire.received.qsize, lambda size: size >= 1, 60) + await asyncio.wait_for( + _answered_while_the_beta_map_is_held(base_url, key, direct_while_held, chat_markers), timeout=45 + ) + assert not held.done(), "Foundry safeguards requests were answered before the beta map loaded" + finally: + release.set() + for served_foundry in await held: + _assert_answered_with_its_own_marker(served_foundry) + + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + beta_by_marker: Final = MappingProxyType( + {marker_of(request): request.headers.get("anthropic-beta") for request in received} + ) + assert {beta_by_marker[call.marker] for call in (*direct, *direct_while_held, *foundry)} == {BETA}, ( + beta_by_marker + ) + assert {beta_by_marker[call.marker] for call in (*plain, *plain_foundry)} == {None}, beta_by_marker + fetches: Final = _beta_map_gets(beta_map_wire.drain()) + assert 1 <= fetches <= 2, fetches diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py new file mode 100644 index 00000000000..0c41efffe73 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -0,0 +1,809 @@ +import itertools +import json +import re +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +BETA: Final = "dangerous-tool-use-2026-09-03" +SAFEGUARDS: Final[list[JsonValue]] = [ + {"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}} +] +ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929" +ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +_CLAUDE_CODE_BETAS: Final = ( + "claude-code-20250219", + "interleaved-thinking-2025-05-14", + "fine-grained-tool-streaming-2025-05-14", + "context-management-2025-06-27", +) +_CLIENT_BETAS: Final = ",".join(_CLAUDE_CODE_BETAS) +_CLIENT_BETAS_WITH_DANGEROUS: Final = ",".join((*_CLAUDE_CODE_BETAS, BETA)) +_MERGED_BETAS: Final = ",".join(sorted((*_CLAUDE_CODE_BETAS, BETA))) +_FOUNDRY_MODEL: Final = "azure_ai/claude-sonnet-4-6" +_FOUNDRY_KEY: Final = "synthetic-foundry-key" +_FOUNDRY_TARGET: Final = "/anthropic/v1/messages" +_FOUNDRY_ARRAY_ERROR: Final = "safeguards: Input should be a valid array" +_FOUNDRY_EXTRA_INPUT_ERROR: Final = "safeguards: Extra inputs are not permitted" +_ANTHROPIC_VERSION: Final = MappingProxyType({"anthropic-version": "2023-06-01"}) +_MESSAGES_KEYS: Final = frozenset({"max_tokens", "messages", "model", "stream"}) +_RESPONSES_KEYS: Final = frozenset({"include", "input", "max_output_tokens", "model"}) +_SDK_EXTRA: Final[dict[str, JsonValue]] = {"safeguards": SAFEGUARDS, "cache": {"no-cache": True}, "num_retries": 0} +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Omitted: + pass + + +OMITTED: Final = _Omitted() + + +@dataclass(frozen=True, slots=True) +class _Observed: + method: str + target: str + marker: str + headers: Mapping[str, str] + body: Mapping[str, JsonValue] + + @property + def beta(self) -> str | None: + return self.headers.get("anthropic-beta") + + @property + def keys(self) -> frozenset[str]: + return frozenset(self.body) + + @property + def header_names(self) -> frozenset[str]: + return frozenset(self.headers) + + +def marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": f"answer marker-{marker}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 5}, + } + + +def stream_frames(marker: str) -> tuple[bytes, ...]: + opening: Final[dict[str, JsonValue]] = { + **_message(marker), + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + } + events: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("message_start", {"message": opening}), + ("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": f"answer marker-{marker}"}}), + ("content_block_stop", {"index": 0}), + ("message_delta", {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 5}}), + ("message_stop", {}), + ) + return tuple( + f"event: {kind}\ndata: {json.dumps({'type': kind, **payload})}\n\n".encode() for kind, payload in events + ) + + +def anthropic_peer(request: Request) -> Reply: + marker: Final = marker_of(request) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=stream_frames(marker)) + return Reply(body=json.dumps(_message(marker)).encode()) + + +def _responses_peer(request: Request) -> Reply: + marker: Final = marker_of(request) + return Reply( + body=json.dumps( + { + "id": f"resp_{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": f"answer marker-{marker}", "annotations": []}], + } + ], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } + ).encode() + ) + + +def _error_peer(status: int, kind: str, message: str) -> Callable[[Request], Reply]: + body: Final = json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode() + return lambda request: Reply(status=status, body=body) + + +def messages_body( + model: str, + marker: str, + *, + safeguards: JsonValue | _Omitted = SAFEGUARDS, + stream: bool = False, + no_cache: bool = True, +) -> dict[str, JsonValue]: + base: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Question marker-{marker}"}], + "num_retries": 0, + **({"stream": True} if stream else {}), + **({"cache": {"no-cache": True}} if no_cache else {}), + } + if isinstance(safeguards, _Omitted): + return base + return {**base, "safeguards": safeguards} + + +def _observe(request: Request) -> _Observed: + return _Observed( + request.method, request.target, marker_of(request), request.headers, _JSON_OBJECT.validate_json(request.body) + ) + + +def _only(wire: Wire) -> _Observed: + (request,) = wire.drain() + return _observe(request) + + +def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response: + headers: Final = {**_ANTHROPIC_VERSION, **({"anthropic-beta": beta} if beta else {})} + return gateway.request("POST", "/v1/messages", body, headers=headers) + + +def _assert_answered(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert response.json()["id"] == f"msg_{marker}", response.text + + +def _assert_streamed(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert f"msg_{marker}" in response.text and f"answer marker-{marker}" in response.text, response.text + assert "event: message_stop" in response.text, response.text + + +def _spend_row(identity: str) -> None: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +def _cache_hit_spend_row(identity: str) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', + (f"{identity}\\_cache\\_hit%",), + ), + lambda found: len(found) == 1, + seconds=70, + ) + assert (rows[0]["status"], str(rows[0]["cache_hit"])) == ("success", "True"), rows + + +def _responses_spend_row(model: str) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s AND call_type = %s', + (model, "aresponses"), + ), + lambda found: len(found) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +def _expect_messages( + observed: _Observed, + *, + target: str, + auth: tuple[str, str] | None, + beta: str | None, + safeguards: JsonValue | _Omitted = SAFEGUARDS, + keys: frozenset[str] = _MESSAGES_KEYS, +) -> None: + assert (observed.method, observed.target) == ("POST", target) + if auth is not None: + credential_matches: Final = observed.headers.get(auth[0]) == auth[1] + assert credential_matches, f"{auth[0]} differs from the deployment credential; headers {observed.header_names}" + assert observed.beta == beta, observed.header_names + if isinstance(safeguards, _Omitted): + assert observed.keys == keys + return + assert observed.keys == keys | {"safeguards"} + assert observed.body["safeguards"] == safeguards + + +def _expect_anthropic_version(observed: _Observed) -> None: + assert observed.headers.get("anthropic-version") == "2023-06-01", observed.header_names + + +@dataclass(frozen=True, slots=True) +class _ShapeCase: + sent: JsonValue | _Omitted + forwarded: JsonValue | _Omitted + client_beta: str | None + expected_beta: str | None + stream: bool = False + + +_ANTHROPIC_SHAPES: Final = MappingProxyType( + { + "D1_safeguards_without_client_header": _ShapeCase(SAFEGUARDS, SAFEGUARDS, None, BETA), + "D2_safeguards_with_claude_code_header": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, _MERGED_BETAS + ), + "D3_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS, _MERGED_BETAS + ), + "D4_empty_safeguards": _ShapeCase([], [], None, BETA), + "D5_null_safeguards": _ShapeCase(None, OMITTED, None, None), + "D6_missing_safeguards": _ShapeCase(OMITTED, OMITTED, None, None), + } +) +_FOUNDRY_SHAPES: Final = MappingProxyType( + { + "D9_streamed_safeguards_with_claude_code_header": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, BETA, stream=True + ), + "D11_empty_safeguards": _ShapeCase([], [], None, BETA), + "D12_null_safeguards": _ShapeCase(None, OMITTED, None, None), + "D13_missing_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase( + OMITTED, OMITTED, _CLIENT_BETAS, None + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_ANTHROPIC_SHAPES)) +def test_anthropic_messages_derive_the_dangerous_tool_use_beta_from_safeguards(gateway: Gateway, row: str) -> None: + case: Final = _ANTHROPIC_SHAPES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + response: Final = _post(gateway, messages_body(model, marker, safeguards=case.sent), beta=case.client_beta) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + beta=case.expected_beta, + safeguards=case.forwarded, + ) + _expect_anthropic_version(observed) + _spend_row(f"msg_{marker}") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _DropCase: + deployment: Callable[[Gateway], Mapping[str, JsonValue]] + target: str + auth: tuple[str, str] | None + keys: frozenset[str] = _MESSAGES_KEYS + + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_MODEL: Final = "claude-opus-5-5" + +_DROP_CASES: Final = MappingProxyType( + { + "anthropic": _DropCase( + deployment=lambda _: {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_key": ANTHROPIC_KEY}, + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + ), + "azure_ai": _DropCase( + deployment=lambda _: {"model": _FOUNDRY_MODEL, "api_key": _FOUNDRY_KEY}, + target=_FOUNDRY_TARGET, + auth=("x-api-key", _FOUNDRY_KEY), + ), + "bedrock_claude_platform": _DropCase( + deployment=lambda _: { + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + target="/v1/messages", + auth=("x-api-key", "synthetic-platform-key"), + ), + "vertex_ai": _DropCase( + deployment=lambda gateway: { + "model": f"vertex_ai/{_VERTEX_MODEL}", + "api_key": None, + "vertex_project": _VERTEX_PROJECT, + "vertex_location": _VERTEX_LOCATION, + "vertex_credentials": service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")), + }, + target=( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + f"/publishers/anthropic/models/{_VERTEX_MODEL}:rawPredict" + ), + auth=None, + keys=frozenset({"anthropic_version", "max_tokens", "messages", "stream"}), + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_DROP_CASES)) +def test_a_deployment_dropping_safeguards_sends_neither_the_field_nor_the_beta(gateway: Gateway, row: str) -> None: + case: Final = _DROP_CASES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=wire.url, additional_drop_params=["safeguards"], **case.deployment(gateway) + ) + response: Final = _post(gateway, messages_body(model, marker)) + _assert_answered(response, marker) + _expect_messages(_only(wire), target=case.target, auth=case.auth, beta=None, safeguards=OMITTED, keys=case.keys) + _spend_row(f"msg_{marker}") + + +def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_its_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", + api_base=wire.url, + api_key=ANTHROPIC_KEY, + additional_drop_params=["context_management"], + ) + body: Final = { + **messages_body(model, marker, safeguards=OMITTED), + "context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}, + } + response: Final = _post(gateway, body) + _assert_answered(response, marker) + _expect_messages( + _only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None, safeguards=OMITTED + ) + _spend_row(f"msg_{marker}") + + +@pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES)) +def test_azure_ai_messages_derive_the_beta_instead_of_forwarding_the_client_header(gateway: Gateway, row: str) -> None: + case: Final = _FOUNDRY_SHAPES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post( + gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta + ) + if case.stream: + _assert_streamed(response, marker) + else: + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, + target=_FOUNDRY_TARGET, + auth=("x-api-key", _FOUNDRY_KEY), + beta=case.expected_beta, + safeguards=case.forwarded, + ) + _expect_anthropic_version(observed) + _spend_row(f"msg_{marker}") + + +def test_anthropic_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + reply: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + extra_body=_SDK_EXTRA, + ) + assert reply.id == f"msg_{marker}", reply + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +async def test_anthropic_sdk_async_stream_safeguards_derive_the_beta_and_log_spend(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) + async with client.messages.with_streaming_response.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + stream=True, + extra_body=_SDK_EXTRA, + ) as streamed: + pieces: Final = [ + event.delta.text + async for event in await streamed.parse() + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ] + assert "".join(pieces) == f"answer marker-{marker}", pieces + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +def test_azure_ai_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + reply: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + extra_body=_SDK_EXTRA, + ) + assert reply.id == f"msg_{marker}", reply + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ProviderCase: + parameters: Mapping[str, JsonValue] + auth: tuple[str, str] | None + expected_beta: str | None + target: str = "/v1/messages" + api_base_suffix: str = "" + model_info: Mapping[str, JsonValue] | None = None + client_beta: str | None = None + keys: frozenset[str] = _MESSAGES_KEYS + extra_headers: Mapping[str, str] = MappingProxyType({}) + body_beta: list[str] | None = None + + +_PROVIDERS: Final = MappingProxyType( + { + "D14_deepseek_with_claude_code_header": _ProviderCase( + parameters={"model": "deepseek/deepseek-chat", "api_key": "synthetic-deepseek-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + target="/anthropic/v1/messages", + auth=("x-api-key", "synthetic-deepseek-key"), + expected_beta=None, + ), + "D15_tencent_with_claude_code_header": _ProviderCase( + parameters={"model": "tencent/deepseek-v4-flash", "api_key": "synthetic-tencent-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-tencent-key"), + expected_beta=None, + ), + "D16_minimax_with_claude_code_header": _ProviderCase( + parameters={"model": "minimax/MiniMax-M2", "api_key": "synthetic-minimax-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-minimax-key"), + expected_beta=None, + ), + "D17_openai_compatible_native_passthrough": _ProviderCase( + parameters={"model": "openai/claude-sonnet-4-5", "api_key": "synthetic-openai-like-key"}, + api_base_suffix="/v1", + model_info={"supported_endpoints": ["/v1/messages"]}, + auth=("authorization", "Bearer synthetic-openai-like-key"), + expected_beta=None, + ), + "D18_sail_with_claude_code_header": _ProviderCase( + parameters={"model": "sail/sail-claude", "api_key": "synthetic-sail-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("authorization", "Bearer synthetic-sail-key"), + expected_beta=None, + ), + "D19_edenai": _ProviderCase( + parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key"}, + auth=("authorization", "Bearer synthetic-edenai-key"), + expected_beta=None, + ), + "D21_bedrock_claude_platform": _ProviderCase( + parameters={ + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + auth=("x-api-key", "synthetic-platform-key"), + expected_beta=BETA, + extra_headers={"anthropic-workspace-id": "ws-synthetic"}, + ), + "D22_bedrock_claude_platform_with_claude_code_header": _ProviderCase( + parameters={ + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-platform-key"), + expected_beta=_MERGED_BETAS, + extra_headers={"anthropic-workspace-id": "ws-synthetic"}, + ), + "D23_bedrock_invoke": _ProviderCase( + parameters={ + "model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + target="/model/anthropic.claude-sonnet-4-5-20250929-v1:0/invoke", + auth=("authorization", "Bearer synthetic-bedrock-bearer"), + expected_beta=None, + keys=frozenset({"anthropic_beta", "anthropic_version", "max_tokens", "messages"}), + body_beta=[BETA], + ), + "D24_bedrock_mantle": _ProviderCase( + parameters={ + "model": "bedrock_mantle/anthropic.claude-sonnet-5-v1:0", + "api_key": "synthetic-mantle-bearer", + "aws_region_name": "us-east-1", + }, + target="/anthropic/v1/messages", + auth=("authorization", "Bearer synthetic-mantle-bearer"), + expected_beta=BETA, + keys=frozenset({"max_tokens", "messages", "model"}), + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_PROVIDERS)) +def test_anthropic_compatible_providers_receive_the_beta_only_where_the_proxy_derives_it( + gateway: Gateway, row: str +) -> None: + case: Final = _PROVIDERS[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model_info=case.model_info, api_base=wire.url + case.api_base_suffix, **case.parameters + ) + response: Final = _post(gateway, messages_body(model, marker), beta=case.client_beta) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages(observed, target=case.target, auth=case.auth, beta=case.expected_beta, keys=case.keys) + for name, value in case.extra_headers.items(): + assert observed.headers.get(name) == value, observed.header_names + assert observed.body.get("anthropic_beta") == case.body_beta, observed.body + _spend_row(f"msg_{marker}") + + +def test_github_copilot_messages_forward_safeguards_without_deriving_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + token_dir: Final = tmp_path / "copilot" + token_dir.mkdir() + with wire_server(anthropic_peer) as wire: + (token_dir / "api-key.json").write_text( + json.dumps({"token": "synthetic-copilot-token", "expires_at": 4102444800, "endpoints": {"api": wire.url}}) + ) + with ( + owned_proxy(gateway, tmp_path, {"GITHUB_COPILOT_TOKEN_DIR": str(token_dir)}) as owned, + owned.scenario() as scenario, + ): + model: Final = scenario.model(model="github_copilot/claude-sonnet-4-5", api_base=wire.url) + response: Final = _post(owned, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, target="/v1/messages", auth=("authorization", "Bearer synthetic-copilot-token"), beta=None + ) + assert observed.headers.get("openai-intent") == "messages-proxy", observed.header_names + _spend_row(f"msg_{marker}") + + +def test_non_native_deployments_bridge_to_the_responses_api_without_safeguards_or_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-openai-key" + ) + response: Final = _post(gateway, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS) + assert response.status_code == 200 and f"answer marker-{marker}" in response.text, response.text + observed: Final = _only(wire) + assert (observed.method, observed.target) == ("POST", "/v1/responses") + assert observed.headers.get("authorization") == "Bearer synthetic-openai-key", observed.header_names + assert observed.beta is None, observed.header_names + assert observed.keys == _RESPONSES_KEYS, observed.keys + _spend_row(str(_JSON_OBJECT.validate_json(response.content)["id"])) + + +def test_chat_completions_forward_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = openai.OpenAI( + base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"Question marker-{marker}"}], extra_body=_SDK_EXTRA + ) + assert completion.choices[0].message.content == f"answer marker-{marker}", completion + _expect_messages( + _only(wire), + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + beta=None, + keys=frozenset({"max_tokens", "messages", "model"}), + ) + _spend_row(completion.id) + + +async def test_responses_stream_forwards_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = openai.AsyncOpenAI( + base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0 + ) + async with client.responses.with_streaming_response.create( + model=model, input=f"Question marker-{marker}", stream=True, extra_body=_SDK_EXTRA + ) as streamed: + events: Final = [event async for event in await streamed.parse()] + deltas: Final = [event.delta for event in events if event.type == "response.output_text.delta"] + assert "".join(deltas) == f"answer marker-{marker}", events + assert any(event.type == "response.completed" for event in events), events + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None) + _responses_spend_row(model) + + +_FOUNDRY_REJECTED_SHAPES: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"S1_int": 1, "S2_empty_string": "", "S3_five_kb_string": "x" * 5000} +) + + +@pytest.mark.parametrize("row", tuple(_FOUNDRY_REJECTED_SHAPES)) +def test_foundry_rejecting_a_non_array_safeguards_value_reaches_the_caller_after_one_attempt( + gateway: Gateway, row: str +) -> None: + value: Final = _FOUNDRY_REJECTED_SHAPES[row] + marker: Final = uuid.uuid4().hex + with ( + wire_server(_error_peer(400, "invalid_request_error", _FOUNDRY_ARRAY_ERROR)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post(gateway, messages_body(model, marker, safeguards=value)) + assert response.status_code == 400 and _FOUNDRY_ARRAY_ERROR in response.text, response.text + _expect_messages( + _only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA, safeguards=value + ) + + +def test_duplicate_safeguards_keys_resolve_to_the_last_value(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + raw: Final = ( + f'{{"model": "{model}", "max_tokens": 16, "num_retries": 0, "cache": {{"no-cache": true}}, ' + f'"messages": [{{"role": "user", "content": "Question marker-{marker}"}}], ' + f'"safeguards": {json.dumps(SAFEGUARDS)}, "safeguards": null}}' + ) + response: Final = gateway.client.post( + "/v1/messages", + content=raw.encode(), + headers={ + "Authorization": f"Bearer {gateway.key}", + "content-type": "application/json", + **_ANTHROPIC_VERSION, + }, + ) + _assert_answered(response, marker) + _expect_messages( + _only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=None, safeguards=OMITTED + ) + _spend_row(f"msg_{marker}") + + +def test_the_same_safeguards_sent_twice_are_two_upstream_calls_with_two_spend_rows(gateway: Gateway) -> None: + markers: Final = (uuid.uuid4().hex, uuid.uuid4().hex) + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + responses: Final = tuple(_post(gateway, messages_body(model, marker)) for marker in markers) + for marker, response in zip(markers, responses, strict=True): + _assert_answered(response, marker) + received: Final = tuple(_observe(request) for request in wire.drain()) + assert len(received) == 2, received + for observed in received: + _expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + assert sorted(item.marker for item in received) == sorted(markers) + for marker in markers: + _spend_row(f"msg_{marker}") + + +_UPSTREAM_ERRORS: Final = MappingProxyType( + { + "S6_unauthorized_deployment": (401, "authentication_error", "invalid x-api-key"), + "S7_foundry_extra_inputs_rejection": (400, "invalid_request_error", _FOUNDRY_EXTRA_INPUT_ERROR), + } +) + + +@pytest.mark.parametrize("row", tuple(_UPSTREAM_ERRORS)) +def test_upstream_errors_reach_the_caller_with_the_beta_sent_once(gateway: Gateway, row: str) -> None: + status, kind, message = _UPSTREAM_ERRORS[row] + marker: Final = uuid.uuid4().hex + with wire_server(_error_peer(status, kind, message)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post(gateway, messages_body(model, marker)) + assert response.status_code == status and message in response.text, response.text + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + + +def _post_noting_whether_upstream_was_called( + gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire +) -> tuple[httpx.Response, bool]: + before: Final = wire.received.qsize() + response: Final = _post(gateway, body) + assert response.status_code == 200, response.text + return response, wire.received.qsize() > before + + +def _until_served_from_cache(gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire) -> tuple[httpx.Response, int]: + deadline: Final = time.monotonic() + 20 + for attempt in itertools.count(1): + response, called_upstream = _post_noting_whether_upstream_was_called(gateway, body, wire) + if not called_upstream: + return response, attempt + assert time.monotonic() < deadline, "the response cache never served the repeated request" + raise AssertionError("unreachable") + + +def test_the_response_cache_serves_the_repeated_request_without_an_upstream_call(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + body: Final = messages_body(model, marker, no_cache=False) + first: Final = _post(gateway, body) + _assert_answered(first, marker) + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + hit, attempts = _until_served_from_cache(gateway, body, wire) + assert hit.json()["id"] == first.json()["id"], hit.text + assert hit.json()["content"] == first.json()["content"], hit.text + _cache_hit_spend_row(f"msg_{marker}") + misses: Final = tuple(_observe(request) for request in wire.drain()) + assert len(misses) == attempts - 1, (attempts, misses) + for observed in misses: + _expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py index dcfd1d19ac2..a448a8cd2da 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py @@ -10,6 +10,7 @@ from anthropic.types import Message from integration._support import claude_code as cc from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.safeguards_rig import DANGEROUS_TOOL_USE_BETA from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -280,7 +281,9 @@ def test_anthropic_messages_preserves_cli_beta_safeguards_and_result(gateway: Ga def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == "/v1/messages", request.target assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers - assert sorted(request.headers["anthropic-beta"].split(",")) == sorted(beta.split(",")), request.headers + assert sorted(request.headers["anthropic-beta"].split(",")) == sorted( + (*beta.split(","), DANGEROUS_TOOL_USE_BETA) + ), request.headers expected_body: Final = { "model": "claude-sonnet-4-6", "max_tokens": 32, diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py index c9da24484e2..57ab6e0e745 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py @@ -4,28 +4,35 @@ with no verdicts.""" from __future__ import annotations +import json from collections.abc import Callable, Iterator, Mapping +from itertools import chain from typing import Final import pytest -from integration._support.client import object_value +from integration._support.anthropic_refusal import DOCS_STOP_DETAILS +from integration._support.client import list_value, object_value from integration._support.mcp import McpPeer, mcp_peer, tool_calls from integration._support.process import graceful_stop_seconds from integration._support.safeguards_rig import ( BASH_TOOL, CHAT_MODEL, + CLAUDE_MODEL, Responder, Rig, auto_mode_body, available, chat_backend, json_body, + native_message, nonce, request_json, safeguards_rig, + sse_events, stream_results, stream_tool_use_ids, tool_id, + tool_use_block, tool_use_ids, ) from integration._support.wire import Reply, Request @@ -115,3 +122,45 @@ def test_auto_executed_mcp_tool_costs_no_classifier_call(rig: Rig, peer: McpPeer received: Final = rig.received(request_nonce) assert rig.classifier_calls(received) == (), [request.target for request in received] assert len(rig.backend_calls(received)) == 2, [request.target for request in received] + + +def carries_a_tool_result_block(request: Request) -> bool: + messages: Final = list_value(request_json(request).get("messages", [])) + contents: Final = tuple(object_value(message).get("content") for message in messages) + blocks: Final = chain.from_iterable(content for content in contents if isinstance(content, list)) + return any(object_value(block).get("type") == "tool_result" for block in blocks) + + +def claude_mcp_tool_then_refusal(request_nonce: str) -> Responder: + identity: Final = tool_id(request_nonce) + tool_turn: Final = native_message( + identity, + content=(tool_use_block(f"{identity}m", name=f"{MCP_ALIAS}-add", arguments={"a": 2, "b": 3}),), + model="claude-sonnet-4-6", + ) + refusal_turn: Final = { + **native_message(f"{identity}r", content=(), stop_reason="refusal", model="claude-sonnet-4-6"), + "stop_details": DOCS_STOP_DETAILS, + } + + def respond(request: Request) -> Reply: + turn: Final = refusal_turn if carries_a_tool_result_block(request) else tool_turn + return Reply(body=json.dumps(turn).encode()) + + return respond + + +def test_auto_executed_mcp_loop_streams_the_final_turn_refusal_stop_details(rig: Rig, peer: McpPeer) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, claude_mcp_tool_then_refusal(request_nonce)) + peer.drain() + status, raw = rig.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, tools=[dict(AUTO_EXECUTED), dict(BASH_TOOL)])) + assert status == 200, raw + assert executed_tools(peer) == ("add",) + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta" + ) + assert len(deltas) == 1, raw + assert deltas[0]["stop_reason"] == "refusal", raw + assert deltas[0].get("stop_details") == DOCS_STOP_DETAILS, raw + assert len(rig.backend_calls(rig.received(request_nonce))) == 2 diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py new file mode 100644 index 00000000000..99a6dd18d48 --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py @@ -0,0 +1,470 @@ +"""A /v1/messages stream LiteLLM builds from a complete reply (web search interception, a fallback's answer, a guardrail +rewrite) carries the ``stop_details`` and ``safeguard_results`` a real Anthropic stream puts on its ``message_delta``.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Iterator, Mapping, Sequence +from typing import Final +from urllib.parse import urlsplit + +import pytest +from integration._support.client import eventually, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.safeguards_rig import ( + BASH_TOOL, + CLAUDE_MODEL, + CLAUDE_TARGET, + TAVILY_TARGET, + Responder, + Rig, + auto_mode_body, + available, + deployment, + native_message, + native_stream_frames, + nonce, + request_json, + safeguards_rig, + sse_events, + stream_results, + stream_tool_use_ids, + text_block, + tool_id, + tool_use_block, + unavailable, +) +from integration._support.wire import Reply, Request +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +CLAUDE: Final = "anthropic/claude-sonnet-4-6" +PRIMARY_MODEL: Final = "synthesized-primary" +FALLBACK_MODEL: Final = "synthesized-fallback" +REFUSED_MODEL: Final = "synthesized-refused" +MOCK_MODEL: Final = "synthesized-mock" +PIPELINE_MODEL: Final = "synthesized-pipeline" +AZURE_REFUSED_MODEL: Final = "synthesized-azure-refused" +CACHED_FALLBACK_MODEL: Final = "synthesized-cached-fallback" +PRIMARY_TARGET: Final = "/primary/v1/messages" +FALLBACK_TARGET: Final = "/fallback/v1/messages" +REFUSED_TARGET: Final = "/refused/v1/messages" +PIPELINE_TARGET: Final = "/pipeline/v1/messages" +AZURE_REFUSED_TARGET: Final = "/azure/anthropic/v1/messages" +CACHED_FALLBACK_TARGET: Final = "/cached-fallback/v1/messages" +MOCK_ANSWER: Final = "mocked fallback answer" +TOOL_PERMISSION: Final = "synthesized-tool-permission" +PIPELINE_REWRITE: Final = "synthesized-pipeline-rewrite" +SECRET: Final = "SECRET" +MASKED: Final = "[MASKED]" +REWRITE_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " if input_type != 'response':\n" + " return allow()\n" + f" texts = [text.replace({SECRET!r}, {MASKED!r}) for text in inputs.get('texts') or []]\n" + " calls = [\n" + " {'id': call['id'], 'type': 'function', 'function': {'name': call['function']['name'],\n" + f" 'arguments': call['function']['arguments'].replace({SECRET!r}, {MASKED!r})}}}}\n" + " for call in inputs.get('tool_calls') or []\n" + " ]\n" + " return modify(texts=texts, tool_calls=calls)\n" +) +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 1} +READ_TOOL: Final[Mapping[str, JsonValue]] = { + "name": "Read", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, +} +NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +REFUSAL_STOP_DETAILS: Final[Mapping[str, JsonValue]] = { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", +} +_FRAME_BOUNDARY: Final = re.compile(rb"(?<=\n\n)") + +LITELLM_SETTINGS: Final[Mapping[str, JsonValue]] = { + "callbacks": ["websearch_interception"], + "websearch_interception_params": {"enabled_providers": ["anthropic"], "search_tool_name": "tavily-search"}, +} + + +def extra_models(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + deployment(PRIMARY_MODEL, CLAUDE, f"{wire_url}/primary"), + deployment(FALLBACK_MODEL, CLAUDE, f"{wire_url}/fallback"), + deployment(REFUSED_MODEL, CLAUDE, f"{wire_url}/refused"), + deployment(MOCK_MODEL, CLAUDE, f"{wire_url}/mock", mock_response=MOCK_ANSWER), + deployment(PIPELINE_MODEL, CLAUDE, f"{wire_url}/pipeline"), + deployment(AZURE_REFUSED_MODEL, "azure_ai/claude-sonnet-4-6", f"{wire_url}/azure/anthropic"), + deployment(CACHED_FALLBACK_MODEL, CLAUDE, f"{wire_url}/cached-fallback"), + ) + + +def top_level(wire_url: str) -> Mapping[str, JsonValue]: + tavily: Final[Mapping[str, JsonValue]] = { + "search_tool_name": "tavily-search", + "litellm_params": { + "search_provider": "tavily", + "api_key": "synthetic-tavily-key", + "api_base": f"{wire_url}/tavily", + }, + } + tool_permission: Final[Mapping[str, JsonValue]] = { + "guardrail_name": TOOL_PERMISSION, + "litellm_params": { + "guardrail": "tool_permission", + "mode": "post_call", + "default_on": False, + "default_action": "allow", + "on_disallowed_action": "rewrite", + "rules": [{"id": "deny_bash", "tool_name": "^Bash$", "decision": "deny"}], + }, + } + pipeline_rewrite: Final[Mapping[str, JsonValue]] = { + "guardrail_name": PIPELINE_REWRITE, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": REWRITE_CODE, + }, + } + masking_policy: Final[Mapping[str, JsonValue]] = { + "guardrails": {"add": [PIPELINE_REWRITE]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": PIPELINE_REWRITE, "on_pass": "allow", "on_fail": "block"}], + }, + } + return { + "search_tools": [dict(tavily)], + "guardrails": [dict(tool_permission), dict(pipeline_rewrite)], + "policies": {"synthesized-masking": dict(masking_policy)}, + "policy_attachments": [{"policy": "synthesized-masking", "models": [PIPELINE_MODEL]}], + "router_settings": { + "disable_cooldowns": True, + "content_policy_fallbacks": [ + {PRIMARY_MODEL: [FALLBACK_MODEL]}, + {REFUSED_MODEL: [MOCK_MODEL]}, + {AZURE_REFUSED_MODEL: [CACHED_FALLBACK_MODEL]}, + ], + }, + } + + +@pytest.fixture(scope="module") +def synthesizing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("synthesized-streams") + with safeguards_rig( + directory, litellm_settings=LITELLM_SETTINGS, top_level=top_level, extra_models=extra_models + ) as rig: + yield rig + + +def claude_backend(message: Mapping[str, JsonValue], delta_fields: Mapping[str, JsonValue]) -> Responder: + def respond(request: Request) -> Reply: + if request_json(request).get("stream") is True: + return Reply( + content_type="text/event-stream", chunks=native_stream_frames(message, delta_extra=delta_fields) + ) + return Reply(body=json.dumps({**message, **delta_fields}).encode()) + + return respond + + +def by_target(backends: Mapping[str, Responder]) -> Responder: + def respond(request: Request) -> Reply: + return backends[urlsplit(request.target).path](request) + + return respond + + +def refusal(identity: str) -> Mapping[str, JsonValue]: + return native_message(identity, content=(), stop_reason="refusal", model="claude-sonnet-4-6") + + +def answer(identity: str, text: str) -> Mapping[str, JsonValue]: + return native_message(identity, content=(text_block(text),), stop_reason="end_turn", model="claude-sonnet-4-6") + + +def refusal_fields() -> Mapping[str, JsonValue]: + return {"stop_details": dict(REFUSAL_STOP_DETAILS), "safeguard_results": list(available({}))} + + +def search_body(model: str, request_nonce: str) -> Mapping[str, JsonValue]: + return auto_mode_body(model, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for", **NO_CACHE) + + +def message_delta(raw: bytes) -> Mapping[str, JsonValue]: + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta" + ) + assert len(deltas) == 1, raw + return deltas[0] + + +def message_start(raw: bytes) -> Mapping[str, JsonValue]: + starts: Final = tuple( + object_value(event["message"]) for event in sse_events(raw) if event.get("type") == "message_start" + ) + assert len(starts) == 1, raw + return starts[0] + + +def streamed_text(raw: bytes) -> str: + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "content_block_delta" + ) + return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +def delta_frame(raw: bytes) -> bytes: + frames: Final = tuple(frame for frame in _FRAME_BOUNDARY.split(raw) if b'"message_delta"' in frame) + assert len(frames) == 1, raw + return frames[0] + + +def targets(rig: Rig, request_nonce: str) -> tuple[str, ...]: + return tuple(urlsplit(request.target).path for request in rig.received(request_nonce)) + + +def test_intercepted_refusal_stream_delivers_the_upstream_stop_details(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields())) + status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce)) + assert status == 200, raw + assert message_delta(raw) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": dict(REFUSAL_STOP_DETAILS), + "safeguard_results": list(available({})), + }, raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_intercepted_refusal_stream_with_a_content_policy_fallback_delivers_the_refusal_without_falling_back( + synthesizing: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + synthesizing.register( + request_nonce, + by_target( + { + PRIMARY_TARGET: claude_backend(refusal(identity), refusal_fields()), + FALLBACK_TARGET: claude_backend(answer(f"{identity}f", f"fallback answer {request_nonce}"), {}), + } + ), + ) + status, raw = synthesizing.stream(search_body(PRIMARY_MODEL, request_nonce)) + assert status == 200, raw + assert message_delta(raw) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": dict(REFUSAL_STOP_DETAILS), + "safeguard_results": list(available({})), + }, raw + assert targets(synthesizing, request_nonce) == (PRIMARY_TARGET,) + + +def test_tool_permission_rewrite_keeps_only_the_delivered_tool_verdict_with_its_explanation(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + denied: Final = tool_id(request_nonce) + allowed: Final = f"{denied}r" + allowed_verdict: Final[Mapping[str, JsonValue]] = { + "type": "evaluated", + "outcome": "flagged", + "explanation": f"reads a secret {allowed}", + } + verdicts: Final[Mapping[str, JsonValue]] = { + denied: {"type": "evaluated", "outcome": "flagged", "explanation": f"destructive {denied}"}, + allowed: dict(allowed_verdict), + } + message: Final = native_message( + denied, + content=( + text_block("Cleaning up."), + tool_use_block(denied), + tool_use_block(allowed, name="Read", arguments={"path": "/work/repo/.env"}), + ), + model="claude-sonnet-4-6", + ) + synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))})) + body: Final = auto_mode_body( + CLAUDE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], guardrails=[TOOL_PERMISSION], **NO_CACHE + ) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert stream_tool_use_ids(raw) == (allowed,), raw + assert stream_results(raw) == available({allowed: dict(allowed_verdict)}), raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_intercepted_answer_stream_keeps_its_safeguard_results(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + text: Final = f"found it {request_nonce}" + results: Final = list(available({})) + synthesizing.register( + request_nonce, + claude_backend(answer(tool_id(request_nonce), text), {"stop_details": None, "safeguard_results": results}), + ) + status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce)) + assert status == 200, raw + assert streamed_text(raw) == text, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw + received: Final = targets(synthesizing, request_nonce) + assert received == (CLAUDE_TARGET,) and TAVILY_TARGET not in received, received + + +def test_native_refusal_stream_without_a_guardrail_keeps_the_upstream_message_delta_bytes(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + message: Final = native_message( + identity, content=(text_block("Partial."),), stop_reason="refusal", model="claude-sonnet-4-6" + ) + fields: Final = {**refusal_fields(), "safeguard_results": list(unavailable((identity,), "error"))} + sent: Final = native_stream_frames(message, delta_extra=fields) + synthesizing.register(request_nonce, claude_backend(message, fields)) + status, raw = synthesizing.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, **NO_CACHE)) + assert status == 200, raw + assert delta_frame(raw) == delta_frame(b"".join(sent)), raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_native_refusal_falling_back_to_a_mock_reply_renders_a_fresh_message_lifecycle(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields())) + status, raw = synthesizing.stream(auto_mode_body(REFUSED_MODEL, request_nonce, **NO_CACHE)) + assert status == 200, raw + started: Final = message_start(raw) + assert started["stop_reason"] is None, raw + assert started.get("stop_details") is None and "safeguard_results" not in started, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None}, raw + assert streamed_text(raw) == MOCK_ANSWER, raw + assert b'"refusal"' not in raw, raw + assert targets(synthesizing, request_nonce) == (REFUSED_TARGET,) + + +def test_native_refusal_falling_back_to_a_cached_intercepted_reply_streams_its_fields_in_message_delta( + synthesizing: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + text: Final = f"cached answer {request_nonce}" + results: Final = list(available({})) + synthesizing.register( + request_nonce, + by_target( + { + AZURE_REFUSED_TARGET: claude_backend(refusal(identity), refusal_fields()), + CACHED_FALLBACK_TARGET: claude_backend( + answer(f"{identity}c", text), {"stop_details": None, "safeguard_results": results} + ), + } + ), + ) + cacheable: Final = auto_mode_body( + CACHED_FALLBACK_MODEL, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for" + ) + + def prime() -> tuple[str, ...]: + primed: Final = synthesizing.post({**cacheable, "stream": False}) + assert primed.status_code == 200, primed.text + return targets(synthesizing, request_nonce) + + first_prime: Final = prime() + assert first_prime == (CACHED_FALLBACK_TARGET,), first_prime + eventually(prime, lambda received: received == ()) + status, raw = synthesizing.stream({**cacheable, "model": AZURE_REFUSED_MODEL}) + assert status == 200, raw + assert streamed_text(raw) == text, raw + started: Final = message_start(raw) + assert started["stop_details"] is None and "safeguard_results" not in started, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw + assert targets(synthesizing, request_nonce) == (AZURE_REFUSED_TARGET,) + + +def flagged_verdict(explanation: str) -> Mapping[str, JsonValue]: + return {"type": "evaluated", "outcome": "flagged", "explanation": explanation} + + +FLAGGED_WITHOUT_EXPLANATION: Final[Mapping[str, JsonValue]] = {"type": "evaluated", "outcome": "flagged"} + + +def masked_refusal(identity: str) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]: + message: Final = native_message( + identity, content=(text_block(f"Partial {SECRET}."),), stop_reason="refusal", model="claude-sonnet-4-6" + ) + expected: Final[Mapping[str, JsonValue]] = { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {**REFUSAL_STOP_DETAILS, "explanation": None}, + "safeguard_results": list(available({})), + } + return message, refusal_fields(), expected + + +def masked_tool_turn( + identity: str, +) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]: + message: Final = native_message( + identity, + content=(text_block(f"Reading {SECRET}."), tool_use_block(identity, name="Read", arguments={"path": "/work"})), + model="claude-sonnet-4-6", + ) + sent: Final[Mapping[str, JsonValue]] = { + "safeguard_results": list(available({identity: dict(flagged_verdict(f"reads {identity}"))})) + } + expected: Final[Mapping[str, JsonValue]] = { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": list(available({identity: dict(FLAGGED_WITHOUT_EXPLANATION)})), + } + return message, sent, expected + + +MASKED_STREAMS: Final = {"refusal_stop_details": masked_refusal, "verdict_explanation": masked_tool_turn} + + +@pytest.mark.parametrize("shape", tuple(MASKED_STREAMS)) +def test_pipeline_text_rewrite_drops_explanations_from_the_upstream_message_delta( + synthesizing: Rig, shape: str +) -> None: + request_nonce: Final = nonce() + message, sent, expected = MASKED_STREAMS[shape](tool_id(request_nonce)) + synthesizing.register(request_nonce, claude_backend(message, sent)) + body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(READ_TOOL)], **NO_CACHE) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert MASKED in streamed_text(raw) and SECRET not in streamed_text(raw), raw + assert message_delta(raw) == expected, raw + assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,) + + +def test_pipeline_tool_call_rewrite_keeps_only_the_untouched_tool_verdict(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + rewritten: Final = tool_id(request_nonce) + untouched: Final = f"{rewritten}r" + message: Final = native_message( + rewritten, + content=( + tool_use_block(rewritten, arguments={"command": f"cat {SECRET}"}), + tool_use_block(untouched, name="Read", arguments={"path": "/work"}), + ), + model="claude-sonnet-4-6", + ) + verdicts: Final[Mapping[str, JsonValue]] = { + rewritten: dict(flagged_verdict(f"reads a file {rewritten}")), + untouched: dict(flagged_verdict(f"reads a directory {untouched}")), + } + synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))})) + body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], **NO_CACHE) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert stream_tool_use_ids(raw) == (rewritten, untouched), raw + assert SECRET.encode() not in raw and MASKED.encode() in raw, raw + assert stream_results(raw) == available({untouched: dict(FLAGGED_WITHOUT_EXPLANATION)}), raw + assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,) diff --git a/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py new file mode 100644 index 00000000000..c55bd181ecb --- /dev/null +++ b/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py @@ -0,0 +1,1015 @@ +"""A post_call guardrail on the Bedrock Converse streaming passthrough streams frame by frame (LIT-9515). + +The caller is boto3 ``ConverseStream`` (or raw HTTP) against ``/bedrock/model//converse-stream``. A +guardrail that scans responses through ``apply_guardrail`` runs on the stream the way it runs on +``/v1/chat/completions``, ``/v1/messages`` and ``/v1/responses``: frames reach the caller as the upstream sends +them, every fifth frame and the end of the stream are scanned, and a block ends the stream with an AWS +exception frame. ``streaming_buffer_until_moderated: true`` still holds the stream and applies masks per content +block, and a caller leaving before the first held frame stops the proxy reading the upstream. Noma v2 scans carry +``stream`` and, on the last one, the finished ``response``, as chat completions sends them; other guardrails see the +request as before.""" + +from __future__ import annotations + +import json +import os +import threading +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Final + +import boto3 +import httpx +import pytest +import yaml +from botocore.config import Config +from botocore.eventstream import EventStreamBuffer +from botocore.exceptions import ClientError, EventStreamError +from integration._support.budget import case_timeout_seconds, import_seconds +from integration._support.client import Gateway, JsonValue, Scenario, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server + +MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +ROUTER_MODEL: Final = "converse-relay" +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +SAMPLED: Final = "relay-sampled" +BUFFERED: Final = "relay-buffered" +CODE: Final = "relay-code" +FAILING: Final = "relay-failing" +FAILING_HELD: Final = "relay-failing-held" +END_ONLY: Final = "relay-end-only" +OUTAGE: Final = "relay-outage" +STACKED: Final = "relay-stacked" +RELEASE_ON_SCAN: Final = "relay-release-on-scan" +REPLACING: Final = "relay-replacing" +REPLACING_THEN_HELD: Final = "relay-replacing-then-held" +PARALLEL: Final = "relay-parallel" +PARALLEL_FIRST: Final = "relay-parallel-first" +HELD_END_ONLY: Final = "relay-held-end-only" +NOMA: Final = "relay-noma" +ALICE: Final = "relay-alice" +NOMA_THEN_ALICE: Final = "relay-noma-then-alice" +CODE_BLOCK_REASON: Final = "code says no" +BLOCK_REASON: Final = "Violated guardrail policy" +MASKED_WORD: Final = "Paris" +MASK: Final = "[CITY]" +BLOCKED_WORD: Final = "FORBIDDEN" +REPLACED_WORD: Final = "CLASSIFIED" +REPLACEMENT: Final = "That answer was withheld" + + +def event_frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload, separators=(",", ":")).encode(), + ) + + +def converse_frames(*blocks: Sequence[str]) -> tuple[bytes, ...]: + """A Converse stream with one text content block per entry of ``blocks``, one delta frame per string.""" + return ( + event_frame("messageStart", {"role": "assistant"}), + *( + frame + for index, deltas in enumerate(blocks) + for frame in ( + *(event_frame("contentBlockDelta", {"contentBlockIndex": index, "delta": {"text": t}}) for t in deltas), + event_frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ), + event_frame("messageStop", {"stopReason": "end_turn"}), + event_frame("metadata", {"usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, "metrics": {}}), + ) + + +@dataclass(frozen=True, slots=True) +class Frame: + message_type: str + kind: str + payload: Mapping[str, JsonValue] + + +def decode(raw: bytes) -> tuple[Frame, ...]: + buffer: Final = EventStreamBuffer() + buffer.add_data(raw) + return tuple( + Frame( + str(message.headers[":message-type"]), + str(message.headers.get(":event-type") or message.headers.get(":exception-type")), + json.loads(message.payload), + ) + for message in buffer + ) + + +def block_texts(frames: Sequence[Frame]) -> dict[int, str]: + texts: Final[dict[int, str]] = {} + for frame in frames: + if frame.kind != "contentBlockDelta": + continue + index = frame.payload["contentBlockIndex"] + delta = frame.payload["delta"] + assert isinstance(index, int) and isinstance(delta, dict), frame + texts[index] = texts.get(index, "") + str(delta["text"]) + return texts + + +def verdict(request: Request) -> Reply: + texts: Final = [str(text) for text in json.loads(request.body)["texts"]] + if any(BLOCKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode()) + if any(MASKED_WORD in text for text in texts): + masked: Final = [text.replace(MASKED_WORD, MASK) for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": masked}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def parallel_verdict(request: Request) -> Reply: + """A block-only check that rejects ``MASKED_WORD`` it sees unmasked.""" + texts: Final = [str(text) for text in json.loads(request.body)["texts"]] + if any(MASKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def bedrock_texts(request: Request) -> tuple[str, ...]: + return tuple(str(item["text"]["text"]) for item in json.loads(request.body)["content"]) + + +def bedrock_verdict(request: Request) -> Reply: + """ApplyGuardrail anonymizing ``MASKED_WORD``, as a Bedrock sensitive-information policy does.""" + texts: Final = bedrock_texts(request) + if not any(MASKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + pii: Final = {"match": MASKED_WORD, "type": "ADDRESS", "action": "ANONYMIZED"} + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": text.replace(MASKED_WORD, MASK)} for text in texts], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [pii]}}], + } + ).encode() + ) + + +def bedrock_replacing_verdict(request: Request) -> Reply: + """ApplyGuardrail blocking ``REPLACED_WORD`` with its blocked message, as a Bedrock word policy does.""" + if not any(REPLACED_WORD in text for text in bedrock_texts(request)): + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + word: Final = {"match": REPLACED_WORD, "action": "BLOCKED"} + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": REPLACEMENT}], + "assessments": [{"wordPolicy": {"customWords": [word]}}], + } + ).encode() + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + scenario: Scenario + provider: Wire + sink: Wire + scripts: dict[str, Reply] + keys: Mapping[str, str] + outage: threading.Event + + def scans(self, guardrail: str, marker: str) -> tuple[tuple[str, ...], ...]: + return tuple( + tuple(json.loads(request.body)["texts"]) + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def bedrock_scans(self, guardrail: str, marker: str) -> tuple[tuple[str, ...], ...]: + return tuple( + bedrock_texts(request) + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def forwarded_request_data(self, marker: str) -> dict[str, tuple[Mapping[str, JsonValue], ...]]: + """The ``request_data`` each response scan forwarded to Noma v2 and Alice, by guardrail, in scan order.""" + forwarded: Final[dict[str, list[Mapping[str, JsonValue]]]] = {NOMA: [], ALICE: []} + for request in self.sink.drain(): + guardrail = request.target.split("/")[1] + if guardrail in forwarded and marker.encode() in request.body: + body = json.loads(request.body) + if body["input_type"] == "response": + forwarded[guardrail].append(body["request_data"]) + return {guardrail: tuple(scans) for guardrail, scans in forwarded.items()} + + def converse_stream(self, guardrail: str, marker: str) -> StreamOutcome: + client: Final = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + endpoint_url=str(self.gateway.client.base_url).rstrip("/") + "/bedrock", + aws_access_key_id="AKIASYNTHETICRELAY", + aws_secret_access_key="synthetic-relay-secret", + config=Config(read_timeout=case_timeout_seconds(), retries={"max_attempts": 1}), + ) + key: Final = self.keys[guardrail] + + def authenticate(request: object, **_: object) -> None: + getattr(request, "headers")["x-litellm-api-key"] = key + + client.meta.events.register("before-send.bedrock-runtime.*", authenticate) + try: + response: Final = client.converse_stream( + modelId=ROUTER_MODEL, messages=[{"role": "user", "content": [{"text": marker}]}] + ) + except ClientError as error: + return StreamOutcome((), None, None, str(error)) + events: Final[list[str]] = [] + try: + for event in response["stream"]: + events.extend(event) + except EventStreamError as error: + return StreamOutcome( + tuple(events), error.response["Error"]["Code"], error.response["Error"]["Message"], None + ) + return StreamOutcome(tuple(events), None, None, None) + + +@dataclass(frozen=True, slots=True) +class StreamOutcome: + events: tuple[str, ...] + error_code: str | None + error_message: str | None + rejected_before_stream: str | None + + +def _guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + scripts: Final[dict[str, Reply]] = {} + outage: Final = threading.Event() + + def provider(request: Request) -> Reply: + assert request.target in (f"/model/{MODEL_ID}/converse-stream", f"/model/{MODEL_ID}/converse"), request.target + marker: Final = json.loads(request.body)["messages"][0]["content"][0]["text"] + return scripts[marker] + + def vendor(request: Request) -> Reply: + failing: Final = request.target.startswith((f"/{FAILING}/", f"/{FAILING_HELD}/")) + if failing or (outage.is_set() and request.target.startswith(f"/{OUTAGE}/")): + return Reply(status=503, body=b'{"error":"guardrail vendor unavailable"}') + if request.target.startswith(f"/{RELEASE_ON_SCAN}/"): + return bedrock_verdict(request) + if request.target.startswith(f"/{REPLACING}/"): + return bedrock_replacing_verdict(request) + if request.target.startswith(f"/{PARALLEL}/"): + return parallel_verdict(request) + if request.target.startswith(f"/{NOMA}/"): + return Reply(body=json.dumps({"action": "NONE"}).encode()) + if request.target.startswith(f"/{ALICE}/"): + return Reply(body=json.dumps({"verdict": "ALLOW"}).encode()) + return verdict(request) + + with httpx.Client(base_url=os.environ["INTEGRATION_PROXY_URL"], timeout=30, trust_env=False) as root_client: + root: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("converse-stream-relay") + with wire_server(provider) as upstream, wire_server(vendor) as sink: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + _guardrail(PARALLEL, sink, run_in_parallel=True), + _guardrail(SAMPLED, sink), + { + "guardrail_name": REPLACING, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": "relay-replacing", + "guardrailVersion": "1", + "aws_access_key_id": "AKIASYNTHETICGUARD", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": f"{sink.url}/{REPLACING}", + "disable_exception_on_block": True, + "streaming_buffer_until_moderated": False, + }, + }, + _guardrail(BUFFERED, sink, streaming_buffer_until_moderated=True), + _guardrail(FAILING, sink), + _guardrail(FAILING_HELD, sink, streaming_buffer_until_moderated=True), + _guardrail(END_ONLY, sink, streaming_end_of_stream_only=True), + _guardrail( + HELD_END_ONLY, sink, streaming_buffer_until_moderated=True, streaming_end_of_stream_only=True + ), + { + "guardrail_name": NOMA, + "litellm_params": { + "guardrail": "noma_v2", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{NOMA}", + "api_key": "synthetic-noma-key", + }, + }, + { + "guardrail_name": ALICE, + "litellm_params": { + "guardrail": "alice", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{ALICE}", + "api_key": "synthetic-alice-key", + }, + }, + _guardrail(OUTAGE, sink), + { + "guardrail_name": RELEASE_ON_SCAN, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": "relay-release-on-scan", + "guardrailVersion": "1", + "aws_access_key_id": "AKIASYNTHETICGUARD", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": f"{sink.url}/{RELEASE_ON_SCAN}", + "streaming_buffer_release_on_scan": True, + }, + }, + { + "guardrail_name": CODE, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " for text in inputs['texts']:\n" + f" if '{BLOCKED_WORD}' in text:\n" + f" return block('{CODE_BLOCK_REASON}')\n" + " return allow()\n" + ), + }, + }, + ] + config["model_list"] = [ + { + "model_name": ROUTER_MODEL, + "litellm_params": { + "model": f"bedrock/{MODEL_ID}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICRELAY", + "aws_secret_access_key": "synthetic-relay-secret", + "aws_region_name": "us-east-1", + }, + } + ] + config_path: Final = directory / "converse-stream-relay.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(root, directory, {}, config=config_path, workers=2) as owned: + with owned.gateway.scenario() as scenario: + single: Final = ( + SAMPLED, + BUFFERED, + CODE, + FAILING, + FAILING_HELD, + END_ONLY, + OUTAGE, + RELEASE_ON_SCAN, + NOMA, + ) + yield Rig( + owned.gateway, + scenario, + upstream, + sink, + scripts, + { + **{name: scenario.key(guardrails=[name]) for name in single}, + STACKED: scenario.key(guardrails=[SAMPLED, BUFFERED]), + REPLACING_THEN_HELD: scenario.key(guardrails=[REPLACING, BUFFERED]), + PARALLEL_FIRST: scenario.key(guardrails=[PARALLEL, BUFFERED]), + NOMA_THEN_ALICE: scenario.key(guardrails=[NOMA, ALICE]), + }, + outage, + ) + + +def _stream_path() -> str: + return f"/bedrock/model/{ROUTER_MODEL}/converse-stream" + + +def _stream_body(marker: str) -> dict[str, JsonValue]: + return {"messages": [{"role": "user", "content": [{"text": marker}]}]} + + +def _stream(rig: Rig, guardrail: str, marker: str) -> tuple[int, str, bytes]: + status, content_type, body, _ = _stream_with_call_id(rig, guardrail, marker) + return status, content_type, body + + +def _stream_with_call_id(rig: Rig, guardrail: str, marker: str) -> tuple[int, str, bytes, str]: + with rig.gateway.client.stream( + "POST", + _stream_path(), + json=_stream_body(marker), + headers={"Authorization": f"Bearer {rig.keys[guardrail]}"}, + timeout=case_timeout_seconds(), + ) as response: + return ( + response.status_code, + response.headers.get("content-type", ""), + response.read(), + response.headers.get("x-litellm-call-id", ""), + ) + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT status, api_base, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s ' + "OR metadata->>'litellm_call_id'=%s ORDER BY status", + (call_id, call_id, call_id), + ) + + +def _spend_rows_for_key_alias(alias: str) -> list[dict[str, JsonValue]]: + return read_rows( + "SELECT status, prompt_tokens FROM \"LiteLLM_SpendLogs\" WHERE metadata->>'user_api_key_alias'=%s", + (alias,), + ) + + +def _stream_releasing_upstream( + rig: Rig, marker: str, first_chunk_length: int, gate: threading.Event +) -> tuple[int, bool, bytes]: + """Stream ``marker`` through the sampled guardrail and release the upstream held after its first transport chunk + once that chunk has reached the caller; report whether the upstream was still held at that point.""" + release_held_upstream: Final = threading.Timer(import_seconds(), gate.set) + release_held_upstream.start() + try: + with rig.gateway.client.stream( + "POST", + _stream_path(), + json=_stream_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + chunks: Final = response.iter_raw() + received: Final = bytearray() + for chunk in chunks: + received.extend(chunk) + if len(received) >= first_chunk_length: + break + upstream_still_held: Final = not gate.is_set() + gate.set() + received.extend(b"".join(chunks)) + return response.status_code, upstream_still_held, bytes(received) + finally: + release_held_upstream.cancel() + gate.set() + + +def test_first_frames_reach_the_caller_while_the_upstream_is_still_streaming(rig: Rig) -> None: + gate: Final = threading.Event() + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Hello {marker} ", "world")) + first_transport_chunk: Final = frames[0] + frames[1] + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=(first_transport_chunk, *frames[2:]), gate_after_first=gate + ) + status, upstream_still_held, received = _stream_releasing_upstream(rig, marker, len(first_transport_chunk), gate) + assert status == 200, received + assert upstream_still_held, "the first frames reached the caller only after the upstream finished its stream" + assert received == b"".join(frames), decode(received) + assert rig.scans(SAMPLED, marker) == ((f"Hello {marker} world",),) + + +def test_block_mid_stream_reaches_boto3_as_an_event_stream_exception(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=converse_frames((f"The secret {marker} is ", BLOCKED_WORD, ".")) + ) + outcome: Final = rig.converse_stream(SAMPLED, marker) + assert outcome.rejected_before_stream is None, outcome + assert outcome.events[:2] == ("messageStart", "contentBlockDelta"), outcome + assert "messageStop" not in outcome.events, outcome + assert (outcome.error_code, outcome.error_message) == ("validationException", BLOCK_REASON), outcome + + +def test_block_raised_mid_stream_logs_the_relayed_spend_and_the_failure(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=converse_frames((f"The secret {marker} is ", BLOCKED_WORD, *(f" word{index}" for index in range(30)))), + pause_between_chunks=0.05, + ) + status, content_type, body, call_id = _stream_with_call_id(rig, CODE, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert frames[0].kind == "messageStart", frames + assert frames[-1].message_type == "exception" and frames[-1].kind == "validationException", frames + assert CODE_BLOCK_REASON in str(frames[-1].payload["message"]), frames + assert call_id, "the relayed stream carries no x-litellm-call-id" + rows: Final = eventually(lambda: _spend_rows(call_id), lambda rows: len(rows) >= 2, seconds=70) + assert [row["status"] for row in rows] == ["failure", "success"], rows + assert rows[1]["api_base"], rows + assert isinstance(rows[1]["prompt_tokens"], int) and rows[1]["prompt_tokens"] > 0, rows + + +def test_clean_stream_reaches_boto3_unchanged(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=converse_frames((f"Fine {marker} ", "answer"))) + outcome: Final = rig.converse_stream(SAMPLED, marker) + assert outcome == StreamOutcome( + ( + "messageStart", + "contentBlockDelta", + "contentBlockDelta", + "contentBlockStop", + "messageStop", + "metadata", + ), + None, + None, + None, + ), outcome + assert rig.scans(SAMPLED, marker) == ((f"Fine {marker} answer",),) + + +def test_sampled_stream_forwards_frames_as_sent_without_applying_masks_like_chat(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"{MASKED_WORD} {marker} ", "is the capital")) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, SAMPLED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + assert rig.scans(SAMPLED, marker) == ((f"{MASKED_WORD} {marker} is the capital",),) + + +def test_a_frame_whose_payload_is_not_utf8_is_relayed_as_sent_and_the_rest_still_scanned(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + clean: Final = converse_frames((f"Before {marker} ", "after")) + garbled: Final = aws_event_stream_frame( + {":event-type": "contentBlockDelta", ":content-type": "application/json", ":message-type": "event"}, + b'{"contentBlockIndex":0,"delta":{"text":"\xff"}}', + ) + frames: Final = (*clean[:2], garbled, *clean[2:]) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, SAMPLED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), body[-300:] + assert set(rig.scans(SAMPLED, marker)) == {(f"Before {marker} after",)} + + +def test_buffered_stream_applies_masks_per_content_block(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=converse_frames((f"{marker} {MASKED_WORD[:3]}", f"{MASKED_WORD[3:]} is big"), ("I like ", MASKED_WORD)), + ) + status, content_type, body = _stream(rig, BUFFERED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert [frame.kind for frame in frames] == [ + "messageStart", + "contentBlockDelta", + "contentBlockDelta", + "contentBlockStop", + "contentBlockDelta", + "contentBlockDelta", + "contentBlockStop", + "messageStop", + "metadata", + ], frames + assert block_texts(frames) == {0: f"{marker} {MASK} is big", 1: f"I like {MASK}"}, frames + assert rig.scans(BUFFERED, marker) == ((f"{marker} {MASKED_WORD} is big", f"I like {MASKED_WORD}"),) + + +def test_buffered_stream_ending_in_an_upstream_exception_still_applies_masks(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + upstream_error: Final = aws_event_stream_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": "Model stream error"}).encode(), + ) + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=(*converse_frames((f"{marker} ", MASKED_WORD))[:3], upstream_error) + ) + status, content_type, body = _stream(rig, BUFFERED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert block_texts(frames) == {0: f"{marker} {MASK}"}, frames + assert frames[-1] == Frame("exception", "modelStreamErrorException", {"message": "Model stream error"}), frames + + +def test_buffered_stream_block_returns_an_error_before_any_frame(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=converse_frames((f"The secret {marker} is ", BLOCKED_WORD)) + ) + outcome: Final = rig.converse_stream(BUFFERED, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the ConverseStream operation: {BLOCK_REASON}" + ), outcome + + +def _spend_statuses(call_id: str, at_least: int) -> list[JsonValue]: + rows: Final = eventually(lambda: _spend_rows(call_id), lambda rows: len(rows) >= at_least, seconds=70) + return [row["status"] for row in rows] + + +def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None: + last: Final = frames[-1] if frames else None + if last is None or last.message_type != "exception": + return None + return last.kind, str(last.payload.get("message")) + + +def test_upstream_closing_mid_stream_ends_the_relayed_stream_with_an_exception_frame(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + gate: Final = threading.Event() + frames: Final = converse_frames((f"Cut {marker} ", "short", " here")) + first_transport_chunk: Final = frames[0] + frames[1] + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=(first_transport_chunk, *frames[2:]), gate_after_first=gate, abort_after=1 + ) + status, upstream_still_held, body = _stream_releasing_upstream(rig, marker, len(first_transport_chunk), gate) + assert status == 200, body + assert upstream_still_held, "the first frames reached the caller only after the upstream closed" + assert body.startswith(first_transport_chunk), decode(body) + relayed: Final = decode(body) + exception: Final = _exception(relayed) + assert exception is not None and exception[0] == "internalServerException", relayed + + +def test_guardrail_vendor_failing_mid_stream_fails_closed_with_an_exception_frame(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Scan {marker} ", "one", " two", " three", " four", " five")) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, FAILING, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body.startswith(b"".join(frames[:4])), decode(body) + relayed: Final = decode(body) + assert "messageStop" not in [frame.kind for frame in relayed], relayed + exception: Final = _exception(relayed) + assert exception is not None and exception[0] == "internalServerException", relayed + + +def test_held_stream_with_the_guardrail_vendor_failing_returns_an_error_before_any_frame(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=converse_frames((f"Held {marker} ", "text"))) + status, content_type, body = _stream(rig, FAILING_HELD, marker) + assert status == 500 and content_type.startswith("application/json"), (status, content_type, body) + + +def test_end_of_stream_only_relays_every_content_frame_then_blocks(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Late {marker} ", "one", " two", " three", " four", f" {BLOCKED_WORD}")) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, END_ONLY, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + content_frames: Final = b"".join(frames[:-2]) + assert body.startswith(content_frames), decode(body) + assert decode(body[len(content_frames) :]) == ( + Frame("exception", "validationException", {"message": BLOCK_REASON}), + ), decode(body) + assert rig.scans(END_ONLY, marker) == ((f"Late {marker} one two three four {BLOCKED_WORD}",),) + + +def test_end_of_stream_only_scans_a_clean_stream_once_and_relays_it_unchanged(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Fine {marker} ", "one", " two", " three", " four", " five")) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, _, body = _stream(rig, END_ONLY, marker) + assert status == 200 and body == b"".join(frames), decode(body) + assert rig.scans(END_ONLY, marker) == ((f"Fine {marker} one two three four five",),) + + +def test_held_stream_released_on_scan_stays_buffered_and_applies_masks(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=converse_frames((f"{marker} ", MASKED_WORD, " is", " the", " capital", " of", " France")), + ) + status, content_type, body = _stream(rig, RELEASE_ON_SCAN, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert block_texts(decode(body)) == {0: f"{marker} {MASK} is the capital of France"}, decode(body) + assert rig.bedrock_scans(RELEASE_ON_SCAN, marker) == ((f"{marker} {MASKED_WORD} is the capital of France",),) + + +def test_sampled_and_held_guardrails_on_one_key_hold_the_stream_and_mask_it(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=converse_frames((f"{MASKED_WORD} {marker} ", "is", " the", " capital")) + ) + status, content_type, body = _stream(rig, STACKED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert block_texts(decode(body)) == {0: f"{MASK} {marker} is the capital"}, decode(body) + scanned_by: Final = { + request.target.split("/")[1] for request in rig.sink.drain() if marker.encode() in request.body + } + assert scanned_by == {SAMPLED, BUFFERED}, scanned_by + + +def test_held_guardrail_still_masks_after_a_sampled_guardrail_replaces_the_rest_of_the_answer(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=converse_frames((f"{MASKED_WORD} {marker} ", "is", f" {REPLACED_WORD}")) + ) + status, content_type, body = _stream(rig, REPLACING_THEN_HELD, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert block_texts(frames) == {0: f"{MASK} {marker} is {REPLACED_WORD}", 1: REPLACEMENT}, frames + assert frames[-1] == Frame("event", "messageStop", {"stopReason": "guardrail_intervened"}), frames + + +def test_parallel_guardrail_listed_before_a_masker_still_scans_the_masked_text(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=converse_frames((f"{MASKED_WORD} {marker} ", "is the capital")) + ) + status, content_type, body = _stream(rig, PARALLEL_FIRST, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert block_texts(decode(body)) == {0: f"{MASK} {marker} is the capital"}, decode(body) + assert rig.scans(PARALLEL, marker) == ((f"{MASK} {marker} is the capital",),) + + +def test_client_disconnect_stops_the_proxy_reading_the_upstream(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Long {marker} ", *(f" word{index}" for index in range(60)))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames, pause_between_chunks=0.1) + while not rig.provider.disconnected.empty(): + rig.provider.disconnected.get_nowait() + with rig.gateway.client.stream( + "POST", + _stream_path(), + json=_stream_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + assert response.status_code == 200 + received: Final = bytearray() + for chunk in response.iter_raw(): + received.extend(chunk) + if len(received) >= len(b"".join(frames[:3])): + break + target: Final = f"/model/{MODEL_ID}/converse-stream" + closed: Final = eventually( + lambda: rig.provider.disconnected.qsize(), lambda count: count > 0, seconds=4, return_last_on_timeout=True + ) + assert closed, "the proxy kept reading the upstream after the caller went away" + assert rig.provider.disconnected.get_nowait() == target + + +def test_client_disconnect_logs_prompt_tokens_the_upstream_never_reported(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Long {marker} ", *(f" word{index}" for index in range(60)))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames, pause_between_chunks=0.1) + with rig.gateway.client.stream( + "POST", + _stream_path(), + json=_stream_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + assert response.status_code == 200 + call_id: Final = response.headers.get("x-litellm-call-id", "") + received: Final = bytearray() + for chunk in response.iter_raw(): + received.extend(chunk) + if len(received) >= len(b"".join(frames[:3])): + break + assert call_id, "the relayed stream carries no x-litellm-call-id" + rows: Final = eventually(lambda: _spend_rows(call_id), lambda rows: len(rows) >= 1, seconds=70) + assert [row["status"] for row in rows] == ["success"], rows + assert isinstance(rows[0]["prompt_tokens"], int) and rows[0]["prompt_tokens"] > 0, rows + + +def test_client_leaving_a_held_stream_before_its_first_frame_stops_the_proxy_reading_the_upstream(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Held {marker} ", *(f" word{index}" for index in range(80)))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames, pause_between_chunks=0.1) + alias: Final = f"held-{marker}" + key: Final = rig.scenario.key(guardrails=[HELD_END_ONLY], key_alias=alias) + while not rig.provider.disconnected.empty(): + rig.provider.disconnected.get_nowait() + with pytest.raises(httpx.ReadTimeout): + rig.gateway.client.post( + _stream_path(), + json=_stream_body(marker), + headers={"Authorization": f"Bearer {key}"}, + timeout=httpx.Timeout(case_timeout_seconds(), read=1.5), + ) + closed: Final = eventually( + lambda: rig.provider.disconnected.qsize(), lambda count: count > 0, seconds=3, return_last_on_timeout=True + ) + assert closed, "the proxy kept reading the upstream after the caller left a held stream" + assert rig.provider.disconnected.get_nowait() == f"/model/{MODEL_ID}/converse-stream" + rows: Final = eventually(lambda: _spend_rows_for_key_alias(alias), lambda rows: len(rows) >= 1, seconds=70) + assert [row["status"] for row in rows] == ["success"], rows + assert isinstance(rows[0]["prompt_tokens"], int) and rows[0]["prompt_tokens"] > 0, rows + + +def _follow_up_status_after_leaving_a_held_stream(rig: Rig, marker: str) -> int: + """Leave a held stream of ``marker`` before its first frame on a key with one parallel slot, then report the status + of the next request on that key once the slot is free, or the last status seen when it never frees.""" + key: Final = rig.scenario.key(guardrails=[HELD_END_ONLY], max_parallel_requests=1) + follow_up: Final = f"relay-{uuid.uuid4().hex}" + rig.scripts[follow_up] = Reply(content_type=EVENT_STREAM, chunks=converse_frames((f"Next {follow_up} ", " done"))) + + def follow_up_status() -> int: + return rig.gateway.client.post( + _stream_path(), + json=_stream_body(follow_up), + headers={"Authorization": f"Bearer {key}"}, + timeout=case_timeout_seconds(), + ).status_code + + with pytest.raises(httpx.ReadTimeout): + rig.gateway.client.post( + _stream_path(), + json=_stream_body(marker), + headers={"Authorization": f"Bearer {key}"}, + timeout=httpx.Timeout(case_timeout_seconds(), read=1.5), + ) + return eventually(follow_up_status, lambda status: status == 200, seconds=5, return_last_on_timeout=True) + + +def test_client_leaving_a_held_stream_before_bedrock_sends_a_byte_frees_its_parallel_request_slot(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + first_byte: Final = threading.Event() + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=converse_frames((f"Held {marker} ", " done")), gate_before_first=first_byte + ) + try: + status: Final = _follow_up_status_after_leaving_a_held_stream(rig, marker) + finally: + first_byte.set() + assert status == 200, "the request that left before Bedrock answered still holds the key's only parallel slot" + + +def test_client_leaving_a_held_stream_after_part_of_the_first_frame_frees_its_parallel_request_slot(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + frames: Final = converse_frames((f"Held {marker} ", " done")) + rest_of_stream: Final = threading.Event() + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=(frames[0][:12], frames[0][12:], *frames[1:]), + gate_after_first=rest_of_stream, + ) + try: + status: Final = _follow_up_status_after_leaving_a_held_stream(rig, marker) + finally: + rest_of_stream.set() + assert status == 200, "the request that left with part of Bedrock's first frame still holds the key's only slot" + + +@pytest.mark.parametrize( + "word_count", [12, 9], ids=["a-sampled-scan-sees-the-end", "the-last-sampled-scan-runs-before-the-end"] +) +def test_noma_scans_carry_stream_and_only_the_last_carries_the_finished_response(rig: Rig, word_count: int) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + words: Final = (f"Noma {marker} ", *(f" word{index}" for index in range(word_count - 1))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=converse_frames(words)) + status, _, body = _stream(rig, NOMA, marker) + assert status == 200 and _exception(decode(body)) is None, decode(body) + scans: Final = rig.forwarded_request_data(marker)[NOMA] + assert len(scans) >= 2, scans + assert [scan.get("stream") for scan in scans] == [True] * len(scans), scans + assert [scan.get("response") is not None for scan in scans] == [False] * (len(scans) - 1) + [True], scans + assert scans[-1]["response"] == { + "output": {"message": {"role": "assistant", "content": [{"text": "".join(words)}]}}, + "stopReason": "end_turn", + }, scans[-1] + + +def test_noma_gets_one_scan_of_a_reply_that_only_calls_a_tool_without_input(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + tool_use: Final = {"toolUseId": "tooluse_relay", "name": "get_time", "type": "tool_use"} + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=( + event_frame("messageStart", {"role": "assistant"}), + event_frame("contentBlockStart", {"contentBlockIndex": 0, "start": {"toolUse": tool_use}}), + event_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": ""}}}), + event_frame("contentBlockStop", {"contentBlockIndex": 0}), + event_frame("messageStop", {"stopReason": "tool_use"}), + event_frame("metadata", {"usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, "metrics": {}}), + ), + ) + status, _, body = _stream(rig, NOMA, marker) + assert status == 200 and _exception(decode(body)) is None, decode(body) + finished: Final = { + "output": {"message": {"role": "assistant", "content": [{"toolUse": {**tool_use, "input": {}}}]}}, + "stopReason": "tool_use", + } + scans: Final = rig.forwarded_request_data(marker)[NOMA] + assert [(scan.get("stream"), scan.get("response")) for scan in scans] == [(True, finished)], scans + + +def test_guardrails_after_noma_see_the_streamed_request_unchanged(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + words: Final = (f"Alice {marker} ", *(f" word{index}" for index in range(11))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=converse_frames(words)) + status, _, body = _stream(rig, NOMA_THEN_ALICE, marker) + assert status == 200 and _exception(decode(body)) is None, decode(body) + scans: Final = rig.forwarded_request_data(marker)[ALICE] + assert len(scans) >= 2, scans + assert [(scan.get("stream"), scan.get("response")) for scan in scans] == [(None, None)] * len(scans), scans + + +def test_non_streaming_converse_block_is_unchanged(rig: Rig) -> None: + marker: Final = f"relay-{uuid.uuid4().hex}" + reply: Final = { + "output": {"message": {"role": "assistant", "content": [{"text": f"The secret {marker} is {BLOCKED_WORD}"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, + "metrics": {"latencyMs": 10}, + } + rig.scripts[marker] = Reply(body=json.dumps(reply).encode()) + response: Final = rig.gateway.client.post( + f"/bedrock/model/{ROUTER_MODEL}/converse", + json=_stream_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 400, response.text + assert BLOCK_REASON in response.text, response.text + + +def test_concurrent_streams_each_get_their_own_verdict_and_spend_rows(rig: Rig) -> None: + markers: Final = [f"relay-{uuid.uuid4().hex}" for _ in range(16)] + blocked: Final = frozenset(markers[::2]) + for marker in markers: + words = (f"Burst {marker} ", "one", " two", BLOCKED_WORD if marker in blocked else " three", " four", " five") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=converse_frames(words), pause_between_chunks=0.02) + with ThreadPoolExecutor(max_workers=len(markers)) as pool: + outcomes: Final = dict(zip(markers, pool.map(partial(_stream_with_call_id, rig, SAMPLED), markers))) + for marker, (status, _, body, call_id) in outcomes.items(): + assert status == 200 and call_id, (marker, status, body) + frames = decode(body) + if marker in blocked: + assert _exception(frames) == ("validationException", BLOCK_REASON), (marker, frames) + else: + assert _exception(frames) is None and frames[-1].kind == "metadata", (marker, frames) + for marker, (_, _, _, call_id) in outcomes.items(): + expected = ["failure", "success"] if marker in blocked else ["success"] + assert _spend_statuses(call_id, len(expected)) == expected, marker + + +def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig: Rig) -> None: + def burst(count: int) -> list[tuple[str, tuple[int, str, bytes, str]]]: + markers = [f"relay-{uuid.uuid4().hex}" for _ in range(count)] + for marker in markers: + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=converse_frames((f"Outage {marker} ", "one", " two", " three", " four")), + ) + with ThreadPoolExecutor(max_workers=count) as pool: + return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers))) + + rig.outage.set() + try: + during: Final = burst(8) + finally: + rig.outage.clear() + after: Final = burst(8) + for marker, (status, _, body, _) in during: + exception = _exception(decode(body)) + assert status == 200 and exception is not None and exception[0] == "internalServerException", (marker, body) + for marker, (status, _, body, _) in after: + assert status == 200 and _exception(decode(body)) is None, (marker, body) + assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body)) + for marker, (_, _, _, call_id) in (*during, *after): + assert _spend_statuses(call_id, 1).count("success") == 1, marker diff --git a/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py new file mode 100644 index 00000000000..3ef73f613fe --- /dev/null +++ b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py @@ -0,0 +1,870 @@ +"""A post_call guardrail scans a Claude model's Bedrock invoke-with-response-stream passthrough as it relays (LIT-9618). + +The caller is boto3 ``InvokeModelWithResponseStream`` (or raw HTTP) against +``/bedrock/model//invoke-with-response-stream``. Each ``chunk`` frame carries one base64 Anthropic Messages +stream event, and a guardrail that scans responses through ``apply_guardrail`` runs on those events the way it runs +on ``/bedrock/model//converse-stream`` and ``/v1/messages``: frames reach the caller as the upstream sends +them, every fifth frame and the end of the stream are scanned, a block ends the stream with an AWS exception frame, +and a held stream gets its masks written back into re-encoded ``chunk`` frames. A guardrail the relay cannot run on +frames, one released on scan or run in parallel, gets the whole stream held and scanned once as converse-stream does. +Streams of other model families keep going out as received""" + +from __future__ import annotations + +import base64 +import json +import os +import threading +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Final +from urllib.parse import unquote + +import boto3 +import httpx +import pytest +import yaml +from botocore.config import Config +from botocore.eventstream import EventStreamBuffer +from botocore.exceptions import ClientError, EventStreamError +from integration._support.budget import case_timeout_seconds, import_seconds +from integration._support.client import Gateway, JsonValue, Scenario, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server + +CLAUDE_ID: Final = "anthropic.claude-sonnet-5-v1:0" +LLAMA_ID: Final = "meta.llama4-maverick-17b-instruct-v1:0" +PROFILE_ID: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/relayprofile" +CLAUDE: Final = "invoke-relay" +LLAMA: Final = "invoke-relay-llama" +PROFILE: Final = "invoke-relay-profile" +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +SAMPLED: Final = "invoke-sampled" +HELD: Final = "invoke-held" +END_ONLY: Final = "invoke-end-only" +CODE: Final = "invoke-code" +BEDROCK: Final = "invoke-bedrock" +BEDROCK_MESSAGE: Final = "invoke-bedrock-message" +UNGUARDED: Final = "invoke-unguarded" +OUTAGE: Final = "invoke-outage" +RELEASE_ON_SCAN: Final = "invoke-release-on-scan" +PARALLEL: Final = "invoke-parallel" +NOMA: Final = "invoke-noma" +NOMA_HELD: Final = "invoke-noma-held" +CODE_BLOCK_REASON: Final = "code says no" +BLOCK_REASON: Final = "Violated guardrail policy" +BEDROCK_BLOCK_MESSAGE: Final = "Sorry, the model cannot answer this question." +MASKED_WORD: Final = "Paris" +MASK: Final = "[CITY]" +BLOCKED_WORD: Final = "FORBIDDEN" + + +def chunk_frame(event: Mapping[str, JsonValue]) -> bytes: + encoded: Final = base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, + json.dumps({"bytes": encoded, "p": "abcdefghij"}, separators=(",", ":")).encode(), + ) + + +def claude_frames(*deltas: str) -> tuple[bytes, ...]: + """A Claude invoke stream with one text content block, one ``text_delta`` event per string.""" + return ( + chunk_frame( + { + "type": "message_start", + "message": { + "id": "msg_bdrk_relay", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + } + ), + chunk_frame({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + *( + chunk_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": delta}}) + for delta in deltas + ), + chunk_frame({"type": "content_block_stop", "index": 0}), + chunk_frame( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + } + ), + chunk_frame( + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 9, + "outputTokenCount": 4, + "invocationLatency": 120, + "firstByteLatency": 40, + }, + } + ), + ) + + +def llama_frames(*generations: str) -> tuple[bytes, ...]: + return tuple( + chunk_frame({"generation": text, "prompt_token_count": None, "generation_token_count": 1, "stop_reason": None}) + for text in generations + ) + + +@dataclass(frozen=True, slots=True) +class Frame: + message_type: str + kind: str + payload: Mapping[str, JsonValue] + + +def _frame(headers: Mapping[str, str], payload: bytes) -> Frame: + body: Final = json.loads(payload) + if headers.get(":event-type") != "chunk": + return Frame(headers[":message-type"], headers.get(":event-type") or headers[":exception-type"], body) + event: Final = json.loads(base64.b64decode(body["bytes"])) + return Frame("event", str(event.get("type")), event) + + +def decode(raw: bytes) -> tuple[Frame, ...]: + buffer: Final = EventStreamBuffer() + buffer.add_data(raw) + return tuple(_frame(message.headers, message.payload) for message in buffer) + + +def streamed_text(frames: Sequence[Frame]) -> str: + return "".join( + str(delta["text"]) + for frame in frames + if frame.kind == "content_block_delta" and isinstance(delta := frame.payload["delta"], dict) + ) + + +def verdict(request: Request) -> Reply: + body: Final = json.loads(request.body) + texts: Final = [str(text) for text in body["texts"]] + if any(BLOCKED_WORD in text for text in texts) or BLOCKED_WORD in json.dumps(body.get("tool_calls")): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode()) + if any(MASKED_WORD in text for text in texts): + masked: Final = [text.replace(MASKED_WORD, MASK) for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": masked}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def bedrock_texts(request: Request) -> tuple[str, ...]: + return tuple(str(item["text"]["text"]) for item in json.loads(request.body)["content"]) + + +def bedrock_verdict(request: Request) -> Reply: + """ApplyGuardrail that anonymizes ``MASKED_WORD`` and blocks ``BLOCKED_WORD``, as Bedrock policies do.""" + texts: Final = bedrock_texts(request) + if any(BLOCKED_WORD in text for text in texts): + word: Final = {"match": BLOCKED_WORD, "action": "BLOCKED"} + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": BEDROCK_BLOCK_MESSAGE}], + "assessments": [{"wordPolicy": {"customWords": [word]}}], + } + ).encode() + ) + if not any(MASKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + pii: Final = {"match": MASKED_WORD, "type": "ADDRESS", "action": "ANONYMIZED"} + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": text.replace(MASKED_WORD, MASK)} for text in texts], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [pii]}}], + } + ).encode() + ) + + +@dataclass(frozen=True, slots=True) +class StreamOutcome: + events: tuple[str, ...] + text: str + error_code: str | None + error_message: str | None + rejected_before_stream: str | None + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + scenario: Scenario + provider: Wire + sink: Wire + scripts: dict[str, Reply] + keys: Mapping[str, str] + outage: threading.Event + + def scans(self, guardrail: str, marker: str) -> tuple[tuple[str, ...], ...]: + return tuple( + tuple(json.loads(request.body)["texts"]) + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def bedrock_scans(self, guardrail: str, marker: str) -> tuple[tuple[str, ...], ...]: + return tuple( + bedrock_texts(request) + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def guardrail_request_data(self, guardrail: str, marker: str) -> tuple[Mapping[str, JsonValue], ...]: + """The request data each Noma scan of the marker's stream carried, in the order the scans ran.""" + return tuple( + json.loads(request.body)["request_data"] + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def invoke_stream(self, guardrail: str, marker: str, model: str = CLAUDE) -> StreamOutcome: + client: Final = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + endpoint_url=str(self.gateway.client.base_url).rstrip("/") + "/bedrock", + aws_access_key_id="AKIASYNTHETICRELAY", + aws_secret_access_key="synthetic-relay-secret", + config=Config(read_timeout=case_timeout_seconds(), retries={"max_attempts": 1}), + ) + key: Final = self.keys[guardrail] + + def authenticate(request: object, **_: object) -> None: + getattr(request, "headers")["x-litellm-api-key"] = key + + client.meta.events.register("before-send.bedrock-runtime.*", authenticate) + try: + response: Final = client.invoke_model_with_response_stream( + modelId=model, body=json.dumps(_invoke_body(marker)) + ) + except ClientError as error: + return StreamOutcome((), "", None, None, str(error)) + events: Final[list[Mapping[str, JsonValue]]] = [] + error: EventStreamError | None = None + try: + for event in response["body"]: + events.append(json.loads(event["chunk"]["bytes"])) + except EventStreamError as raised: + error = raised + frames: Final = tuple(Frame("event", str(event["type"]), event) for event in events) + return StreamOutcome( + tuple(frame.kind for frame in frames), + streamed_text(frames), + None if error is None else error.response["Error"]["Code"], + None if error is None else error.response["Error"]["Message"], + None, + ) + + +def _guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +def _bedrock_guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": name, + "guardrailVersion": "1", + "aws_access_key_id": "AKIASYNTHETICGUARD", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": f"{sink.url}/{name}", + **settings, + }, + } + + +def _noma_guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "noma_v2", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +def _deployment(name: str, model_id: str, upstream: Wire) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"bedrock/{model_id}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICRELAY", + "aws_secret_access_key": "synthetic-relay-secret", + "aws_region_name": "us-east-1", + }, + } + + +def _invoke_body(marker: str) -> dict[str, JsonValue]: + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}], + } + + +def _marker(request: Request) -> str: + body: Final = json.loads(request.body) + if "prompt" in body: + return str(body["prompt"]) + return str(body["messages"][0]["content"][0]["text"]) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + scripts: Final[dict[str, Reply]] = {} + outage: Final = threading.Event() + + def provider(request: Request) -> Reply: + assert unquote(request.target) in ( + f"/model/{CLAUDE_ID}/invoke-with-response-stream", + f"/model/{CLAUDE_ID}/invoke", + f"/model/{LLAMA_ID}/invoke-with-response-stream", + f"/model/{PROFILE_ID}/invoke-with-response-stream", + ), request.target + return scripts[_marker(request)] + + def vendor(request: Request) -> Reply: + if outage.is_set() and request.target.startswith(f"/{OUTAGE}/"): + return Reply(status=503, body=b'{"error":"guardrail vendor unavailable"}') + if request.target.startswith((f"/{BEDROCK}/", f"/{BEDROCK_MESSAGE}/", f"/{RELEASE_ON_SCAN}/")): + return bedrock_verdict(request) + if request.target.startswith((f"/{NOMA}/", f"/{NOMA_HELD}/")): + return Reply(body=json.dumps({"action": "NONE"}).encode()) + return verdict(request) + + with httpx.Client(base_url=os.environ["INTEGRATION_PROXY_URL"], timeout=30, trust_env=False) as root_client: + root: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("invoke-stream-relay") + with wire_server(provider) as upstream, wire_server(vendor) as sink: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + _guardrail(SAMPLED, sink), + _guardrail(HELD, sink, streaming_buffer_until_moderated=True), + _guardrail(END_ONLY, sink, streaming_end_of_stream_only=True), + _guardrail(OUTAGE, sink), + _guardrail(PARALLEL, sink, run_in_parallel=True), + _noma_guardrail(NOMA, sink, streaming_buffer_until_moderated=False), + _noma_guardrail(NOMA_HELD, sink, streaming_buffer_until_moderated=True), + _bedrock_guardrail(BEDROCK, sink), + _bedrock_guardrail(BEDROCK_MESSAGE, sink, disable_exception_on_block=True), + _bedrock_guardrail(RELEASE_ON_SCAN, sink, streaming_buffer_release_on_scan=True), + { + "guardrail_name": CODE, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " for text in inputs['texts']:\n" + f" if '{BLOCKED_WORD}' in text:\n" + f" return block('{CODE_BLOCK_REASON}')\n" + " return allow()\n" + ), + }, + }, + ] + config["model_list"] = [ + _deployment(CLAUDE, CLAUDE_ID, upstream), + _deployment(LLAMA, LLAMA_ID, upstream), + _deployment(PROFILE, PROFILE_ID, upstream), + ] + config_path: Final = directory / "invoke-stream-relay.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(root, directory, {}, config=config_path, workers=2) as owned: + with owned.gateway.scenario() as scenario: + guarded: Final = ( + SAMPLED, + HELD, + END_ONLY, + CODE, + BEDROCK, + BEDROCK_MESSAGE, + OUTAGE, + RELEASE_ON_SCAN, + PARALLEL, + NOMA, + NOMA_HELD, + ) + yield Rig( + owned.gateway, + scenario, + upstream, + sink, + scripts, + { + **{name: scenario.key(guardrails=[name]) for name in guarded}, + UNGUARDED: scenario.key(), + }, + outage, + ) + + +def _stream_with_call_id(rig: Rig, guardrail: str, marker: str, model: str = CLAUDE) -> tuple[int, str, bytes, str]: + with rig.gateway.client.stream( + "POST", + f"/bedrock/model/{model}/invoke-with-response-stream", + json={"prompt": marker} if model == LLAMA else _invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[guardrail]}"}, + timeout=case_timeout_seconds(), + ) as response: + return ( + response.status_code, + response.headers.get("content-type", ""), + response.read(), + response.headers.get("x-litellm-call-id", ""), + ) + + +def _stream(rig: Rig, guardrail: str, marker: str, model: str = CLAUDE) -> tuple[int, str, bytes]: + status, content_type, body, _ = _stream_with_call_id(rig, guardrail, marker, model) + return status, content_type, body + + +def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None: + last: Final = frames[-1] if frames else None + if last is None or last.message_type != "exception": + return None + return last.kind, str(last.payload.get("message")) + + +def _spend_statuses(call_id: str, at_least: int) -> list[JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s ' + "OR metadata->>'litellm_call_id'=%s ORDER BY status", + (call_id, call_id, call_id), + ), + lambda rows: len(rows) >= at_least, + seconds=70, + ) + return [row["status"] for row in rows] + + +def test_first_frames_reach_the_caller_while_the_upstream_is_still_streaming_and_get_scanned(rig: Rig) -> None: + gate: Final = threading.Event() + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Hello {marker} ", "world") + first_transport_chunk: Final = frames[0] + frames[1] + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=(first_transport_chunk, *frames[2:]), gate_after_first=gate + ) + release_held_upstream: Final = threading.Timer(import_seconds(), gate.set) + release_held_upstream.start() + try: + with rig.gateway.client.stream( + "POST", + f"/bedrock/model/{CLAUDE}/invoke-with-response-stream", + json=_invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + chunks: Final = response.iter_raw() + received: Final = bytearray() + for chunk in chunks: + received.extend(chunk) + if len(received) >= len(first_transport_chunk): + break + upstream_still_held: Final = not gate.is_set() + gate.set() + received.extend(b"".join(chunks)) + status: Final = response.status_code + finally: + release_held_upstream.cancel() + gate.set() + assert status == 200, received + assert upstream_still_held, "the first frames reached the caller only after the upstream finished its stream" + assert bytes(received) == b"".join(frames), decode(bytes(received)) + assert rig.scans(SAMPLED, marker) == ((f"Hello {marker} world",),) + + +def test_block_mid_stream_reaches_boto3_as_an_event_stream_exception(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD, ".") + ) + outcome: Final = rig.invoke_stream(SAMPLED, marker) + assert outcome.rejected_before_stream is None, outcome + assert outcome.events[:2] == ("message_start", "content_block_start"), outcome + assert "message_stop" not in outcome.events, outcome + assert (outcome.error_code, outcome.error_message) == ("validationException", BLOCK_REASON), outcome + + +def test_claude_behind_an_application_inference_profile_is_scanned_by_its_request_body(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD, ".") + ) + outcome: Final = rig.invoke_stream(SAMPLED, marker, model=PROFILE) + assert outcome.rejected_before_stream is None, outcome + assert "message_stop" not in outcome.events, outcome + assert (outcome.error_code, outcome.error_message) == ("validationException", BLOCK_REASON), outcome + + +def test_clean_stream_reaches_boto3_unchanged_after_its_scan(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=claude_frames(f"Fine {marker} ", "answer")) + outcome: Final = rig.invoke_stream(SAMPLED, marker) + assert outcome == StreamOutcome( + ( + "message_start", + "content_block_start", + "content_block_delta", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ), + f"Fine {marker} answer", + None, + None, + None, + ), outcome + assert rig.scans(SAMPLED, marker) == ((f"Fine {marker} answer",),) + + +def test_held_stream_writes_the_guardrail_masks_back_into_the_chunk_frames(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"{marker} {MASKED_WORD[:3]}", f"{MASKED_WORD[3:]} is big") + ) + status, content_type, body = _stream(rig, HELD, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert [frame.kind for frame in frames] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ], frames + assert streamed_text(frames) == f"{marker} {MASK} is big", frames + assert rig.scans(HELD, marker) == ((f"{marker} {MASKED_WORD} is big",),) + + +def test_held_stream_block_returns_an_error_before_any_frame(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD) + ) + outcome: Final = rig.invoke_stream(HELD, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + + +def test_held_stream_the_upstream_breaks_off_still_scans_its_tool_input(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + text_frames: Final = claude_frames(f"Looking up {marker}") + tool_use: Final = {"type": "tool_use", "id": "toolu_relay", "name": "lookup", "input": {}} + upstream_error: Final = aws_event_stream_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": "upstream stream broke"}).encode(), + ) + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=( + *text_frames[:4], + chunk_frame({"type": "content_block_start", "index": 1, "content_block": tool_use}), + chunk_frame( + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": json.dumps({"query": BLOCKED_WORD})}, + } + ), + chunk_frame({"type": "content_block_stop", "index": 1}), + upstream_error, + ), + ) + outcome: Final = rig.invoke_stream(HELD, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + + +def test_bedrock_guardrail_holds_the_stream_and_anonymizes_it(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"{marker} ", MASKED_WORD, " is", " the", " capital") + ) + outcome: Final = rig.invoke_stream(BEDROCK, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.text == f"{marker} {MASK} is the capital", outcome + assert rig.bedrock_scans(BEDROCK, marker) == ((f"{marker} {MASKED_WORD} is the capital",),) + + +def test_bedrock_guardrail_block_message_replaces_a_held_stream(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD) + ) + outcome: Final = rig.invoke_stream(BEDROCK_MESSAGE, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.events[0] == "message_start" and outcome.events[-1] == "message_stop", outcome + assert outcome.text == BEDROCK_BLOCK_MESSAGE, outcome + + +def test_end_of_stream_only_relays_every_content_frame_then_blocks(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Late {marker} ", "one", " two", " three", " four", f" {BLOCKED_WORD}") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, END_ONLY, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + content_frames: Final = b"".join(frames[:-2]) + assert body.startswith(content_frames), decode(body) + assert decode(body[len(content_frames) :]) == ( + Frame("exception", "validationException", {"message": BLOCK_REASON}), + ), decode(body) + assert rig.scans(END_ONLY, marker) == ((f"Late {marker} one two three four {BLOCKED_WORD}",),) + + +def test_block_raised_mid_stream_logs_the_relayed_spend_and_the_failure(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD, *(f" word{index}" for index in range(30))), + pause_between_chunks=0.05, + ) + status, content_type, body, call_id = _stream_with_call_id(rig, CODE, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert frames[0].kind == "message_start", frames + assert _exception(frames) == ("validationException", CODE_BLOCK_REASON), frames + assert call_id, "the relayed stream carries no x-litellm-call-id" + assert _spend_statuses(call_id, 2) == ["failure", "success"] + + +def test_non_claude_invoke_stream_is_relayed_as_received_without_a_scan(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = llama_frames(f"Llama {marker} ", BLOCKED_WORD, " end") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, SAMPLED, marker, model=LLAMA) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + assert rig.scans(SAMPLED, marker) == () + + +def test_invoke_stream_without_a_guardrail_is_relayed_as_received(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Open {marker} ", BLOCKED_WORD, " end") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, UNGUARDED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + + +def test_held_stream_released_on_scan_stays_buffered_and_applies_masks(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"{marker} ", MASKED_WORD, " is", " the", " capital") + ) + outcome: Final = rig.invoke_stream(RELEASE_ON_SCAN, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.events[0] == "message_start" and outcome.events[-1] == "message_stop", outcome + assert outcome.text == f"{marker} {MASK} is the capital", outcome + assert rig.bedrock_scans(RELEASE_ON_SCAN, marker) == ((f"{marker} {MASKED_WORD} is the capital",),) + + +def test_held_stream_released_on_scan_with_a_block_returns_an_error_before_any_event(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"Release {marker} ", BLOCKED_WORD, " end") + ) + outcome: Final = rig.invoke_stream(RELEASE_ON_SCAN, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + assert rig.bedrock_scans(RELEASE_ON_SCAN, marker) == ((f"Release {marker} {BLOCKED_WORD} end",),) + + +def test_parallel_guardrail_scans_the_whole_invoke_stream_before_any_event(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"Parallel {marker} ", BLOCKED_WORD, " end") + ) + outcome: Final = rig.invoke_stream(PARALLEL, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + assert rig.scans(PARALLEL, marker) == ((f"Parallel {marker} {BLOCKED_WORD} end",),) + + +def test_each_sampled_scan_sends_a_streaming_request_with_every_event_relayed_so_far(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Noma {marker} ", *(f" word{index}" for index in range(9))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, NOMA, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + scans: Final = rig.guardrail_request_data(NOMA, marker) + assert len(scans) > 1, scans + assert [scan.get("stream") for scan in scans] == [True] * len(scans), scans + relayed: Final = [len(scan["responses"]) for scan in scans if isinstance(scan.get("responses"), list)] + assert len(relayed) == len(scans) and relayed == sorted(set(relayed)), relayed + + +def test_the_end_of_stream_scan_sends_the_whole_stream_and_its_finish_reason(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Noma {marker} ", "held", " reply", " here") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + outcome: Final = rig.invoke_stream(NOMA_HELD, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.text == f"Noma {marker} held reply here", outcome + scans: Final = rig.guardrail_request_data(NOMA_HELD, marker) + assert [scan.get("stream") for scan in scans] == [True] * len(scans) and scans, scans + scan: Final = scans[-1] + responses: Final = scan.get("responses") + assert isinstance(responses, list) and len(responses) >= len(frames), responses + response: Final = scan.get("response") + assert isinstance(response, dict) and isinstance(choices := response.get("choices"), list), response + assert [choice.get("finish_reason") for choice in choices if isinstance(choice, dict)] == ["stop"], response + + +def test_upstream_exception_mid_stream_reaches_the_caller_after_the_scanned_frames(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + upstream_error: Final = aws_event_stream_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": "upstream stream broke"}).encode(), + ) + frames: Final = claude_frames(f"Partial {marker} ", "answer") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=(*frames[:3], upstream_error)) + status, content_type, body = _stream(rig, SAMPLED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join((*frames[:3], upstream_error)), decode(body) + assert _exception(decode(body)) == ("modelStreamErrorException", "upstream stream broke") + assert rig.scans(SAMPLED, marker) == ((f"Partial {marker} ",),) + + +def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig: Rig) -> None: + def burst(count: int) -> list[tuple[str, tuple[int, str, bytes, str]]]: + markers = [f"invoke-{uuid.uuid4().hex}" for _ in range(count)] + for marker in markers: + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"Outage {marker} ", "one", " two", " three") + ) + with ThreadPoolExecutor(max_workers=count) as pool: + return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers))) + + rig.outage.set() + try: + during: Final = burst(8) + finally: + rig.outage.clear() + after: Final = burst(8) + for marker, (status, _, body, _) in during: + exception = _exception(decode(body)) + assert status == 200 and exception is not None and exception[0] == "internalServerException", (marker, body) + for marker, (status, _, body, _) in after: + assert status == 200 and _exception(decode(body)) is None, (marker, body) + assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body)) + for marker, (_, _, _, call_id) in (*during, *after): + assert _spend_statuses(call_id, 1).count("success") == 1, marker + + +def test_client_disconnect_stops_the_proxy_reading_the_upstream(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Long {marker} ", *(f" word{index}" for index in range(60))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames, pause_between_chunks=0.1) + while not rig.provider.disconnected.empty(): + rig.provider.disconnected.get_nowait() + with rig.gateway.client.stream( + "POST", + f"/bedrock/model/{CLAUDE}/invoke-with-response-stream", + json=_invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + assert response.status_code == 200 + received: Final = bytearray() + for chunk in response.iter_raw(): + received.extend(chunk) + if len(received) >= len(b"".join(frames[:3])): + break + closed: Final = eventually( + lambda: rig.provider.disconnected.qsize(), lambda count: count > 0, seconds=4, return_last_on_timeout=True + ) + assert closed, "the proxy kept reading the upstream after the caller went away" + assert rig.provider.disconnected.get_nowait() == f"/model/{CLAUDE_ID}/invoke-with-response-stream" + + +def test_non_streaming_invoke_block_is_unchanged(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + reply: Final = { + "id": "msg_bdrk_relay", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": f"The secret {marker} is {BLOCKED_WORD}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 4}, + } + rig.scripts[marker] = Reply(body=json.dumps(reply).encode()) + response: Final = rig.gateway.client.post( + f"/bedrock/model/{CLAUDE}/invoke", + json=_invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 400, response.text + assert BLOCK_REASON in response.text, response.text diff --git a/tests/integration/observability/test_callback_delivery.py b/tests/integration/observability/test_callback_delivery.py index 543e60da4a6..e786e655a13 100644 --- a/tests/integration/observability/test_callback_delivery.py +++ b/tests/integration/observability/test_callback_delivery.py @@ -11,7 +11,7 @@ import pytest import yaml from integration._support.client import Gateway, JsonValue, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -268,6 +268,89 @@ def test_streamed_responses_success_callback_carries_provider_apim_request_id(ga }, json.dumps(events[0]["hidden_params"]) +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_non_stream_openai_like_success_callback_carries_provider_headers(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = uuid.uuid4().hex + request_id: Final = "req_" + marker + sink_secret: Final = "synthetic-sink-secret-" + marker + + def upstream(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + marker, + "object": "chat.completion", + "created": 1, + "model": "google.gemma-4-31b", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode(), + headers={"x-amzn-requestid": request_id, "x-ratelimit-remaining-requests": "997"}, + ) + + def sink(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {sink_secret}" + return Reply() + + with wire_server(upstream) as provider, wire_server(sink) as endpoint: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy( + gateway, + tmp_path, + { + "GENERIC_LOGGER_ENDPOINT": endpoint.url, + "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}", + }, + config=path, + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model="bedrock_mantle/google.gemma-4-31b", + api_base=provider.url, + api_key="synthetic-mantle-bearer", + aws_region_name="us-west-2", + ) + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "header callback " + marker}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == marker, response.text + assert len(provider.drain()) == 1 + batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered() -> tuple[dict, ...]: + batches.extend(endpoint.drain()) + return tuple( + event for batch in batches for event in json.loads(batch.body) if event.get("model_group") == model + ) + + events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=10) + assert (events[0]["status"], events[0]["call_type"]) == ("success", "acompletion") + assert events[0]["litellm_call_id"] == response.headers["x-litellm-call-id"] + additional_headers: Final = events[0]["hidden_params"]["additional_headers"] or {} + assert { + name: additional_headers.get(name) + for name in ("llm_provider-x-amzn-requestid", "x_ratelimit_remaining_requests") + } == {"llm_provider-x-amzn-requestid": request_id, "x_ratelimit_remaining_requests": 997}, json.dumps( + events[0]["hidden_params"] + ) + + _RAISING_HOOK: Final = """ from litellm.integrations.custom_logger import CustomLogger diff --git a/tests/integration/observability/test_otel_cache_hit_cost.py b/tests/integration/observability/test_otel_cache_hit_cost.py new file mode 100644 index 00000000000..76c21b529ef --- /dev/null +++ b/tests/integration/observability/test_otel_cache_hit_cost.py @@ -0,0 +1,370 @@ +import json +import os +import re +import uuid +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, list_value, object_value +from integration._support.database import read_rows +from integration._support.otlp_sink import Span, owned_sinks, recorded_spans +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +INPUT_PRICE: Final = 0.001 +OUTPUT_PRICE: Final = 0.002 +PROMPT_TOKENS: Final = 7 +COMPLETION_TOKENS: Final = 2 +CHAT_PRICE: Final = PROMPT_TOKENS * INPUT_PRICE + COMPLETION_TOKENS * OUTPUT_PRICE +EMBEDDING_PRICE: Final = PROMPT_TOKENS * INPUT_PRICE +CLIENT: Final = 3 +COST_HEADERS: Final = ("x-litellm-response-cost", "x-litellm-response-cost-input", "x-litellm-response-cost-output") +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +MARKER: Final = re.compile(rb"cache-hit-cost-[0-9a-f]{32}") + + +def _sse(frames: tuple[Mapping[str, JsonValue], ...]) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(*(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames), b"data: [DONE]\n\n"), + ) + + +def _answer(marker: str) -> str: + return "answer to " + marker + + +def _embedding_value(marker: str) -> float: + return float(int(marker[-12:], 16)) + + +def _chat_reply(stream: bool, marker: str) -> Reply: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + message: Final = {"role": "assistant", "content": _answer(marker)} + usage: Final = { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + if not stream: + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": usage, + } + ).encode() + ) + chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + return _sse( + ( + {**chunk, "choices": [{"index": 0, "delta": message}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": usage}, + ) + ) + + +def _responses_reply(marker: str) -> Reply: + identity: Final = "resp_" + uuid.uuid4().hex + return Reply( + body=json.dumps( + { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_" + identity, + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": _answer(marker), "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + } + ).encode() + ) + + +def _embedding_reply(marker: str) -> Reply: + return Reply( + body=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [_embedding_value(marker), 0.5, 0.75]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": PROMPT_TOKENS, "total_tokens": PROMPT_TOKENS}, + } + ).encode() + ) + + +def _provider(request: Request) -> Reply: + found: Final = MARKER.search(request.body) + assert found is not None, request.body + marker: Final = found.group().decode() + if request.target.endswith("/embeddings"): + return _embedding_reply(marker) + if request.target.endswith("/responses"): + return _responses_reply(marker) + return _chat_reply(JSON_OBJECT.validate_json(request.body).get("stream") is True, marker) + + +@dataclass(frozen=True, slots=True) +class Route: + path: str + stream: bool + price: float + body: Callable[[str, str], Mapping[str, JsonValue]] + cached: Callable[[str], str] = _answer + + +ROUTES: Final[Mapping[str, Route]] = MappingProxyType( + { + "chat": Route( + "/v1/chat/completions", + False, + CHAT_PRICE, + lambda model, marker: {"model": model, "messages": [{"role": "user", "content": marker}]}, + ), + "chat_stream": Route( + "/v1/chat/completions", + True, + CHAT_PRICE, + lambda model, marker: { + "model": model, + "messages": [{"role": "user", "content": marker}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + ), + "messages": Route( + "/v1/messages", + False, + CHAT_PRICE, + lambda model, marker: {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + ), + "responses": Route("/v1/responses", False, CHAT_PRICE, lambda model, marker: {"model": model, "input": marker}), + "embeddings": Route( + "/v1/embeddings", + False, + EMBEDDING_PRICE, + lambda model, marker: {"model": model, "input": marker}, + lambda marker: repr(_embedding_value(marker)), + ), + } +) + + +@dataclass(frozen=True, slots=True) +class Served: + call_id: str + headers: httpx.Headers + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + sink: str + chat_model: str + embedding_model: str + upstream: Wire + otel_v2: bool + + def send(self, route: str, marker: str) -> Served: + chosen: Final = ROUTES[route] + model: Final = self.embedding_model if route == "embeddings" else self.chat_model + with self.proxy.client.stream( + "POST", + chosen.path, + json=dict(chosen.body(model, marker)), + headers={"Authorization": f"Bearer {self.proxy.key}"}, + ) as response: + body: Final = response.read() + assert response.status_code == 200, body + return Served(response.headers["x-litellm-call-id"], response.headers) + + def upstream_calls(self, marker: str) -> int: + return sum(1 for request in self.upstream.drain() if marker.encode() in request.body) + + def cost_span(self, call_id: str) -> Span: + def matching() -> tuple[Span, ...]: + return tuple(span for span in recorded_spans(self.sink)[1] if self._is_cost_span(span, call_id)) + + found: Final = eventually(matching, lambda spans: len(spans) >= 1, seconds=30) + assert len(found) == 1, found + return found[0] + + def _is_cost_span(self, span: Span, call_id: str) -> bool: + if span["attributes"].get("litellm.call_id") != call_id: + return False + if self.otel_v2: + return span["kind"] == CLIENT and "gen_ai.operation.name" in span["attributes"] + return span["name"] == "litellm_request" + + def total_cost(self, span: Span) -> float: + return _number(span["attributes"].get("litellm.cost.total" if self.otel_v2 else "gen_ai.cost.total_cost")) + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float, str)) and not isinstance(value, bool), value + return float(value) + + +def _config(provider: Wire, sink: str, chat_model: str, embedding_model: str) -> dict[str, JsonValue]: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + priced: Final = {"input_cost_per_token": INPUT_PRICE, "output_cost_per_token": OUTPUT_PRICE} + return { + **config, + "model_list": [ + *list_value(config.get("model_list") or []), + { + "model_name": chat_model, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + **priced, + }, + }, + { + "model_name": embedding_model, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + **priced, + }, + }, + ], + "litellm_settings": {**object_value(config["litellm_settings"]), "callbacks": ["otel"]}, + "callback_settings": {"otel": {"exporter": "http/json", "endpoint": sink, "use_simple_processor": True}}, + } + + +@pytest.fixture(scope="module") +def provider() -> Iterator[Wire]: + with wire_server(_provider) as wire: + yield wire + + +def _rig(provider: Wire, directory: Path, otel_v2: bool) -> Iterator[Rig]: + chat_model: Final = "cachecost-chat-" + uuid.uuid4().hex[:12] + embedding_model: Final = "cachecost-embed-" + uuid.uuid4().hex[:12] + with owned_sinks(directory / "sinks") as sinks: + path: Final = directory / f"cache-hit-cost-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(_config(provider, sinks.operator, chat_model, embedding_model))) + v1_exporter: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator} + exporter: Final = {"LITELLM_OTEL_V2": "1"} if otel_v2 else v1_exporter + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + directory, + exporter, + config=path, + remove_environment=() if otel_v2 else ("LITELLM_OTEL_V2",), + workers=2, + ) as owned, + ): + yield Rig(owned.gateway, sinks.operator, chat_model, embedding_model, provider, otel_v2) + + +@pytest.fixture(scope="module") +def v1(provider: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + yield from _rig(provider, tmp_path_factory.mktemp("cache-hit-cost-v1"), otel_v2=False) + + +@pytest.fixture(scope="module") +def v2(provider: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + yield from _rig(provider, tmp_path_factory.mktemp("cache-hit-cost-v2"), otel_v2=True) + + +def _spend_row(call_id: str) -> Mapping[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, cache_hit, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,) + ), + lambda found: len(found) == 1, + seconds=70, + ) + return rows[0] + + +def _breakdown_total(row: Mapping[str, JsonValue]) -> float: + return _number(object_value(object_value(row["metadata"])["cost_breakdown"])["total_cost"]) + + +def _cache_holds(needle: str) -> bool: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as store: + return any(needle.encode() in _bytes(store.get(key)) for key in store.scan_iter(count=1000, _type="STRING")) + + +def _bytes(value: object) -> bytes: + return value if isinstance(value, bytes) else b"" + + +def _miss_then_hit(rig: Rig, route: str) -> tuple[Served, Served]: + marker: Final = "cache-hit-cost-" + uuid.uuid4().hex + miss: Final = rig.send(route, marker) + eventually(lambda: _cache_holds(ROUTES[route].cached(marker)), bool, seconds=30) + hit: Final = rig.send(route, marker) + assert rig.upstream_calls(marker) == 1 + return miss, hit + + +@pytest.mark.parametrize("route", tuple(ROUTES)) +def test_a_v1_cache_hit_exports_zero_cost_and_answers_zero_cost_headers(v1: Rig, route: str) -> None: + price: Final = ROUTES[route].price + miss, hit = _miss_then_hit(v1, route) + + hit_span: Final = v1.cost_span(hit.call_id) + assert v1.total_cost(hit_span) == 0, hit_span + assert hit_span["attributes"].get("litellm.cache_hit") is True, hit_span + assert _number(hit_span["attributes"].get("litellm.saved_cache_cost")) == pytest.approx(price), hit_span + + assert {name: hit.headers.get(name) for name in COST_HEADERS} == dict.fromkeys(COST_HEADERS, "0.0"), hit.headers + + hit_row: Final = _spend_row(hit.call_id) + assert (_number(hit_row["spend"]), hit_row["cache_hit"]) == (0, "True"), hit_row + assert _breakdown_total(hit_row) == 0, hit_row + assert _number(_spend_row(miss.call_id)["spend"]) == pytest.approx(price) + + miss_span: Final = v1.cost_span(miss.call_id) + assert v1.total_cost(miss_span) == pytest.approx(price), miss_span + assert miss_span["attributes"].get("litellm.cache_hit") is False, miss_span + + +@pytest.mark.parametrize("route", tuple(ROUTES)) +def test_a_v2_cache_hit_exports_an_llm_span_at_zero_cost(v2: Rig, route: str) -> None: + price: Final = ROUTES[route].price + miss, hit = _miss_then_hit(v2, route) + + hit_span: Final = v2.cost_span(hit.call_id) + assert v2.total_cost(hit_span) == 0, hit_span + assert hit_span["attributes"].get("litellm.cache_hit") is True, hit_span + assert _number(hit_span["attributes"].get("litellm.saved_cache_cost")) == pytest.approx(price), hit_span + + miss_span: Final = v2.cost_span(miss.call_id) + assert v2.total_cost(miss_span) == pytest.approx(price), miss_span + assert miss_span["attributes"].get("litellm.cache_hit") is False, miss_span diff --git a/tests/integration/observability/test_presidio_overlapping_spans.py b/tests/integration/observability/test_presidio_overlapping_spans.py new file mode 100644 index 00000000000..a2d92b15330 --- /dev/null +++ b/tests/integration/observability/test_presidio_overlapping_spans.py @@ -0,0 +1,251 @@ +import json +import re +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +MODEL: Final = "gpt-6.1-sol" +PHONE: Final = "4251234567" +EMAIL: Final = "john.smith@example.com" +INTERNATIONAL_PHONE: Final = "+33 1 42 68 53 00" +US_PHONE: Final = "+1 (206) 555-0199" +SAME_SPAN: Final = f"Repeat exactly: my number is {PHONE}" +NESTED: Final = f"contact {EMAIL} today" +CONTAINED_FRAGMENT: Final = f"her phone is {INTERNATIONAL_PHONE}." +SEPARATE_PARTS: Final = (f"My number is {PHONE}.", f"Also call {US_PHONE}.") +# Spans, scores and anonymized texts recorded from presidio-analyzer and presidio-anonymizer 2.2.364 on 2026-10-16 +DETECTIONS: Final = MappingProxyType( + { + PHONE: (("PHONE_NUMBER", 0, 10, 0.75), ("US_BANK_NUMBER", 0, 10, 0.05), ("US_DRIVER_LICENSE", 0, 10, 0.01)), + EMAIL: (("EMAIL_ADDRESS", 0, 22, 1.0), ("URL", 0, 7, 0.5), ("URL", 11, 22, 0.5)), + "today": (("DATE_TIME", 0, 5, 0.85),), + INTERNATIONAL_PHONE: (("DATE_TIME", 12, 17, 0.85), ("PHONE_NUMBER", 0, 17, 0.75)), + US_PHONE: (("PHONE_NUMBER", 0, 17, 0.4),), + } +) +ANONYMIZED: Final = MappingProxyType( + { + SAME_SPAN: "Repeat exactly: my number is ", + NESTED: "contact ", + CONTAINED_FRAGMENT: "her phone is .", + SEPARATE_PARTS[0]: "My number is .", + SEPARATE_PARTS[1]: "Also call .", + } +) +CASES: Final = ( + pytest.param(SAME_SPAN, "Repeat exactly: my number is ", id="same_span"), + pytest.param(NESTED, "contact ", id="nested"), + pytest.param(CONTAINED_FRAGMENT, "her phone is .", id="container_outranks_higher_scored_fragment"), +) +STREAM: Final = pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + + +def _detections(text: str, value: str) -> Iterator[dict[str, str | int | float]]: + for match in re.finditer(rf"(? list[dict[str, str | int | float]]: + findings: Final = chain.from_iterable(_detections(text, value) for value in DETECTIONS) + return sorted(findings, key=lambda item: -item["score"]) + + +def _spans(results: list[dict[str, str | int | float]]) -> list[tuple[str | int | float, ...]]: + return sorted(tuple(result[key] for key in ("entity_type", "start", "end", "score")) for result in results) + + +def _analyze(request: Request) -> Reply: + assert request.target == "/analyze", request.target + return Reply(body=json.dumps(_findings(json.loads(request.body)["text"])).encode()) + + +def _anonymize(request: Request) -> Reply: + assert request.target == "/anonymize", request.target + body: Final = json.loads(request.body) + text: Final = body["text"] + if text not in ANONYMIZED: + return Reply(status=500, body=json.dumps({"error": f"no recorded anonymization for {text!r}"}).encode()) + if _spans(body["analyzer_results"]) != _spans(_findings(text)): + return Reply( + status=500, body=json.dumps({"error": f"analyzer_results differ from the detections for {text!r}"}).encode() + ) + return Reply(body=json.dumps({"text": ANONYMIZED[text], "items": []}).encode()) + + +def _texts(messages: list[dict]) -> tuple[str, ...]: + return tuple( + chain.from_iterable( + (message["content"],) + if isinstance(message["content"], str) + else (part["text"] for part in message["content"]) + for message in messages + ) + ) + + +def _chunk(text: str, finish: str | None) -> bytes: + delta: Final = {"role": "assistant", "content": text} if finish is None else {} + payload: Final = { + "id": "chatcmpl-echo", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _echo(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + body: Final = json.loads(request.body) + text: Final = " | ".join(_texts(body["messages"])) + if body.get("stream"): + pieces: Final = tuple(text[index : index + 6] for index in range(0, len(text), 6)) + return Reply( + content_type="text/event-stream", + chunks=(*(_chunk(piece, None) for piece in pieces), _chunk("", "stop"), b"data: [DONE]\n\n"), + ) + reply: Final = { + "id": "chatcmpl-echo", + "object": "chat.completion", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 9, "total_tokens": 18}, + } + return Reply(body=json.dumps(reply).encode()) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + provider: Wire + + +def _config(root: Path, provider: Wire, analyzer: Wire, anonymizer: Wire, scope: Literal["input"] | None) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["cache"] = False + config["model_list"].append( + { + "model_name": MODEL, + "litellm_params": {"model": f"openai/{MODEL}", "api_base": provider.url + "/v1", "api_key": "synthetic"}, + } + ) + guardrail: Final = { + "guardrail": "presidio", + "mode": "pre_call", + "output_parse_pii": True, + "default_on": True, + "presidio_analyzer_api_base": analyzer.url + "/", + "presidio_anonymizer_api_base": anonymizer.url + "/", + } + scoped: Final = guardrail if scope is None else {**guardrail, "presidio_filter_scope": scope} + config["guardrails"] = [{"guardrail_name": "presidio-pii", "litellm_params": scoped}] + path: Final = root / "presidio-overlapping-spans.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def _rig(root: Path, scope: Literal["input"] | None) -> Iterator[Rig]: + with ( + gateway_from_environment() as gateway, + wire_server(_echo) as provider, + wire_server(_analyze) as analyzer, + wire_server(_anonymize) as anonymizer, + owned_proxy_process( + gateway, root, {}, config=_config(root, provider, analyzer, anonymizer, scope), workers=2 + ) as owned, + ): + yield Rig(owned.gateway, provider) + + +@pytest.fixture(scope="module") +def default_scope_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory.mktemp("presidio-default-scope"), None) as rig: + yield rig + + +@pytest.fixture(scope="module") +def input_scope_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory.mktemp("presidio-input-scope"), "input") as rig: + yield rig + + +def _send(rig: Rig, messages: list[dict], stream: bool) -> tuple[int, str]: + rig.provider.drain() + body: Final = {"model": MODEL, "messages": messages, "stream": stream} + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + if not stream: + response: Final = rig.proxy.client.post("/v1/chat/completions", json=body, headers=headers) + if response.status_code != 200: + return response.status_code, response.text + return response.status_code, response.json()["choices"][0]["message"]["content"] + with rig.proxy.client.stream("POST", "/v1/chat/completions", json=body, headers=headers) as streamed: + if streamed.status_code != 200: + return streamed.status_code, streamed.read().decode() + frames: Final = tuple( + json.loads(line.removeprefix("data: ")) + for line in streamed.iter_lines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + return 200, "".join(frame["choices"][0]["delta"].get("content") or "" for frame in frames if frame["choices"]) + + +def _forwarded(rig: Rig) -> tuple[str, ...]: + received: Final = rig.provider.drain() + assert len(received) == 1, received + return _texts(json.loads(received[0].body)["messages"]) + + +@STREAM +@pytest.mark.parametrize(("text", "forwarded"), CASES) +def test_overlapping_detections_reach_the_provider_as_one_token_and_come_back_unmasked( + input_scope_rig: Rig, text: str, forwarded: str, stream: bool +) -> None: + status, replied = _send(input_scope_rig, [{"role": "user", "content": text}], stream) + assert status == 200, replied + assert _forwarded(input_scope_rig) == (forwarded,) + assert replied == text + + +@STREAM +@pytest.mark.parametrize(("text", "forwarded"), CASES) +def test_default_scope_masks_the_cleanly_restored_reply_like_presidio_masks_the_original( + default_scope_rig: Rig, text: str, forwarded: str, stream: bool +) -> None: + status, replied = _send(default_scope_rig, [{"role": "user", "content": text}], stream) + assert status == 200, replied + assert _forwarded(default_scope_rig) == (forwarded,) + assert replied == ANONYMIZED[text] + + +@STREAM +@pytest.mark.parametrize( + "messages", + [ + pytest.param( + [{"role": "user", "content": [{"type": "text", "text": part} for part in SEPARATE_PARTS]}], + id="content_parts", + ), + pytest.param([{"role": "user", "content": part} for part in SEPARATE_PARTS], id="messages"), + ], +) +def test_one_type_in_separate_parts_gets_distinct_tokens_and_each_part_comes_back_as_itself( + input_scope_rig: Rig, messages: list[dict], stream: bool +) -> None: + status, replied = _send(input_scope_rig, messages, stream) + assert status == 200, replied + assert _forwarded(input_scope_rig) == ("My number is .", "Also call .") + assert replied == " | ".join(SEPARATE_PARTS) diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py index d784fecd146..d69866ec43d 100644 --- a/tests/integration/observability/test_presidio_streaming_output.py +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -2,7 +2,7 @@ import json import re import threading import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager from dataclasses import dataclass @@ -10,6 +10,7 @@ from pathlib import Path from typing import Final import httpx +import pytest import yaml from integration._support.client import Gateway, eventually from integration._support.process import OwnedProxy, owned_proxy_process @@ -457,6 +458,96 @@ def test_anthropic_messages_per_request_guardrails_selects_masking(gateway: Gate assert len(rig.analyzer.drain()) == 1 +READ_TOOL_USE_ID: Final = "toolu_presidio_read" +READ_INPUT: Final = {"path": "/work/report.txt"} + + +def anthropic_tool_turn_stream(identity: str, text: str, delta: Mapping[str, object]) -> tuple[bytes, ...]: + opening, text_start, text_delta, text_stop, _, closing = anthropic_stream(identity, text) + tool_use: Final = {"type": "tool_use", "id": READ_TOOL_USE_ID, "name": "Read", "input": {}} + return ( + opening, + text_start, + text_delta, + text_stop, + anthropic_frame("content_block_start", {"type": "content_block_start", "index": 1, "content_block": tool_use}), + anthropic_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": json.dumps(READ_INPUT)}, + }, + ), + anthropic_frame("content_block_stop", {"type": "content_block_stop", "index": 1}), + anthropic_frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 4}}), + closing, + ) + + +def read_verdicts(verdict: Mapping[str, object]) -> Sequence[Mapping[str, object]]: + return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {READ_TOOL_USE_ID: verdict}}}] + + +def anthropic_deltas(received: Received, event_type: str) -> tuple[Mapping[str, object], ...]: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in received.text.split("\n") if line.startswith("data: ") + ) + return tuple(event["delta"] for event in events if event.get("type") == event_type) + + +def anthropic_message_delta(received: Received) -> Mapping[str, object]: + deltas: Final = anthropic_deltas(received, "message_delta") + assert len(deltas) == 1, received.text + return deltas[0] + + +def anthropic_text_deltas(received: Received) -> str: + deltas: Final = anthropic_deltas(received, "content_block_delta") + return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +@pytest.mark.parametrize( + ("delta", "field", "kept"), + [ + ( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": read_verdicts( + {"type": "evaluated", "outcome": "flagged", "explanation": f"{PERSON} asked for the report"} + ), + }, + "safeguard_results", + read_verdicts({"type": "evaluated", "outcome": "flagged"}), + ), + ( + { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": f"declined for {PERSON}"}, + }, + "stop_details", + {"type": "refusal", "category": "cyber", "explanation": None}, + ), + ], + ids=["verdict_explanation", "stop_details_explanation"], +) +def test_anthropic_messages_masked_stream_keeps_upstream_safeguard_fields_without_explanations( + gateway: Gateway, tmp_path: Path, delta: Mapping[str, object], field: str, kept: object +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + provider: Final = anthropic_provider(anthropic_tool_turn_stream(identity, f"{PERSON} designed it.", delta)) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text_deltas(received) == f"{MASK} designed it.", received.text + assert anthropic_message_id(received) == identity, received.text + assert anthropic_message_delta(received).get(field) == kept, received.text + assert PERSON not in received.text, received.text + assert len(rig.upstream.drain()) == 1 + + def openai_provider(identity: str) -> Callable[[Request], Reply]: def provider(request: Request) -> Reply: assert request.target == "/v1/chat/completions", request.target diff --git a/tests/integration/observability/test_stream_error_alert_api_base.py b/tests/integration/observability/test_stream_error_alert_api_base.py new file mode 100644 index 00000000000..449e378dfc7 --- /dev/null +++ b/tests/integration/observability/test_stream_error_alert_api_base.py @@ -0,0 +1,482 @@ +import asyncio +import hashlib +import json +import uuid +from collections.abc import Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import openai +import pytest +import yaml +from anthropic.types import RawContentBlockDeltaEvent, TextDelta +from integration._support.anthropic_sse import DROP_PAUSE, dropping_reply, stream_reply +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_stream import created, delta, failed, frame +from integration._support.responses_ws_peer import PROVIDER_MODEL, Script, events, responses_peer, session +from integration._support.tls import write_self_signed_cert +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.chat import ChatCompletionChunk +from openai.types.responses import ResponseFailedEvent, ResponseTextDeltaEvent +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +AZURE_KEY: Final = "synthetic-azure-key" +API_VERSION: Final = "2025-04-01-preview" +AZURE_CHAT_BACKEND: Final = "azure/gpt-4o-mini" +PARTIAL_TEXT: Final = "partial answer" +ALERT_SECONDS: Final = 40 +NO_RETRY_NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) +EMPTY_MODEL_LIST: Final = Reply(body=b'{"object": "list", "data": []}') +VERTEX_PROJECT: Final = "alert-vertex-project" +VERTEX_LOCATION: Final = "us-east5" +VERTEX_PARTNER_BACKEND: Final = "claude-opus-5-5@default" +VERTEX_PARTNER_URL: Final = ( + f"https://{VERTEX_LOCATION}-aiplatform.googleapis.com/v1/projects/{VERTEX_PROJECT}" + f"/locations/{VERTEX_LOCATION}/publishers/anthropic/models/{VERTEX_PARTNER_BACKEND}" +) +VERTEX_OVERLOADED: Final = Reply( + status=500, body=b'{"error": {"code": 500, "message": "vertex partner overloaded", "status": "INTERNAL"}}' +) + +API_BASES: Final[Mapping[str, Callable[[str, str], str]]] = MappingProxyType( + { + "plain": lambda url, _: f"{url}/gateway", + "code": lambda url, secret: f"{url}/gateway?code={secret}", + "subscription-key": lambda url, secret: f"{url}/gateway?subscription-key={secret}", + "userinfo": lambda url, secret: url.replace("://", f"://alice:{secret}@", 1) + "/gateway", + } +) + + +@dataclass(frozen=True, slots=True) +class AlertingProxy: + gateway: Gateway + alerts: Wire + vertex: Wire + cert: tuple[Path, Path] + + +@dataclass(frozen=True, slots=True) +class Caller: + key: str + key_alias: str + team_alias: str + + +@dataclass(frozen=True, slots=True) +class Outcome: + text: str + failure: str + + +def accept_alert(request: Request) -> Reply: + return Reply(body=b"ok", content_type="text/plain") + + +def overloaded_vertex(request: Request) -> Reply: + return VERTEX_OVERLOADED + + +def alerting_config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": { + **object_value(base["general_settings"]), + "alerting": ["slack"], + "alert_types": ["llm_exceptions"], + }, + } + path: Final = directory / "alerting.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def alerting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[AlertingProxy]: + directory: Final = tmp_path_factory.mktemp("stream-error-alerts") + cert: Final = write_self_signed_cert(directory) + with ( + gateway_from_environment() as shared, + wire_server(accept_alert) as alerts, + wire_server(overloaded_vertex) as vertex, + ): + with owned_proxy( + shared, + directory, + { + "SLACK_WEBHOOK_URL": alerts.url, + "DEFAULT_FLUSH_INTERVAL_SECONDS": "1", + "SSL_CERT_FILE": str(cert[0]), + "VERTEXAI_API_BASE": vertex.url, + }, + config=alerting_config(directory), + ) as gateway: + yield AlertingProxy(gateway, alerts, vertex, cert) + + +def alert_naming(alerts: Wire, marker: str) -> str: + def matching() -> tuple[str, ...]: + texts: Final = (string_value(JSON_OBJECT.validate_json(request.body)["text"]) for request in alerts.drain()) + return tuple(text for text in texts if marker in text) + + return eventually(matching, lambda found: len(found) > 0, seconds=ALERT_SECONDS)[0] + + +def api_base_line(alert: str) -> str | None: + return next((line for line in alert.splitlines() if line.startswith("API Base:")), None) + + +def assert_names_deployment(alert: str, api_base: str, model_group: str, backend: str, caller: Caller) -> None: + assert api_base_line(alert) == f"API Base: `{api_base}`", alert + assert f"model_group: `{model_group}`" in alert, alert + assert f"deployment: `{backend}`" in alert, alert + assert f"Key Name: `{caller.key_alias}`\nTeam: `{caller.team_alias}`" in alert, alert + + +def keyed_caller(scenario: Scenario) -> Caller: + team_alias: Final = f"alert-team-{uuid.uuid4().hex[:12]}" + key_alias: Final = f"alert-key-{uuid.uuid4().hex[:12]}" + team: Final = scenario.team(team_alias=team_alias) + return Caller(scenario.key(key_alias=key_alias, team_id=team), key_alias, team_alias) + + +def azure_chat(scenario: Scenario, api_base: str, *, bridged: bool = False) -> str: + return scenario.model( + model=AZURE_CHAT_BACKEND, + api_base=api_base, + api_key=AZURE_KEY, + api_version=API_VERSION, + num_retries=0, + use_chat_completions_api=bridged, + ) + + +def provider_message(marker: str) -> str: + return f"upstream overloaded {marker}" + + +def chat_chunk(marker: str, content: Mapping[str, JsonValue]) -> bytes: + body: Final = { + "id": f"chatcmpl-{marker}", + "object": "chat.completion.chunk", + "created": 1, + "model": f"gpt-4o-mini-{marker}", + "choices": [{"index": 0, "delta": dict(content), "finish_reason": None}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def first_text(marker: str) -> bytes: + return chat_chunk(marker, {"role": "assistant", "content": PARTIAL_TEXT}) + + +def cut_after_text(marker: str) -> Reply: + return dropping_reply((first_text(marker), chat_chunk(marker, {"content": " never sent"})), abort_after=1) + + +def error_frame_after_text(marker: str) -> Reply: + error: Final = {"error": {"message": provider_message(marker), "type": "server_error", "code": "server_error"}} + return stream_reply((first_text(marker), b"data: " + json.dumps(error).encode() + b"\n\n"), pause=DROP_PAUSE) + + +def rejected_before_stream(marker: str) -> Reply: + return Reply( + status=500, body=json.dumps({"error": {"message": provider_message(marker), "type": "server_error"}}).encode() + ) + + +def chat_upstream(reply: Callable[[str], Reply], marker: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return EMPTY_MODEL_LIST + return reply(marker) + + return respond + + +def items_or_error(open_stream: Callable[[], Iterable[object]], error_type: type[Exception]) -> Iterator[object]: + try: + yield from open_stream() + except error_type as error: + yield error + + +def chunk_text(item: object) -> str: + if not isinstance(item, ChatCompletionChunk): + return "" + return "".join(choice.delta.content or "" for choice in item.choices) + + +def message_text(item: object) -> str: + if isinstance(item, RawContentBlockDeltaEvent) and isinstance(item.delta, TextDelta): + return item.delta.text + return "" + + +def responses_text(item: object) -> str: + return item.delta if isinstance(item, ResponseTextDeltaEvent) else "" + + +def failure_of(item: object) -> str: + if isinstance(item, Exception): + return str(item) + if isinstance(item, ResponseFailedEvent) and item.response.error is not None: + return item.response.error.message + return "" + + +def outcome(items: tuple[object, ...], text_of: Callable[[object], str]) -> Outcome: + return Outcome("".join(text_of(item) for item in items), "".join(failure_of(item) for item in items)) + + +def openai_client(gateway: Gateway, key: str) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=key, max_retries=0, timeout=60) + + +def chat_outcome(gateway: Gateway, key: str, model: str) -> Outcome: + with openai_client(gateway, key) as client: + items: Final = tuple( + items_or_error( + lambda: client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + extra_body=dict(NO_RETRY_NO_CACHE), + ), + openai.APIError, + ) + ) + return outcome(items, chunk_text) + + +def messages_outcome(gateway: Gateway, key: str, model: str) -> Outcome: + with anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=key, max_retries=0, timeout=60) as client: + items: Final = tuple( + items_or_error( + lambda: client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + extra_body=dict(NO_RETRY_NO_CACHE), + ), + anthropic.APIError, + ) + ) + return outcome(items, message_text) + + +def responses_outcome(gateway: Gateway, key: str, model: str) -> Outcome: + with openai_client(gateway, key) as client: + items: Final = tuple( + items_or_error( + lambda: client.responses.create( + model=model, input="stream an answer", stream=True, extra_body=dict(NO_RETRY_NO_CACHE) + ), + openai.APIError, + ) + ) + return outcome(items, responses_text) + + +ENTRYPOINTS: Final[Mapping[str, Callable[[Gateway, str, str], Outcome]]] = MappingProxyType( + {"chat": chat_outcome, "messages": messages_outcome, "responses-bridge": responses_outcome} +) +FAILURES: Final[Mapping[str, Callable[[str], Reply]]] = MappingProxyType( + {"cut": cut_after_text, "error-frame": error_frame_after_text} +) + + +@pytest.mark.parametrize( + ("entrypoint", "failure"), + [("chat", "cut"), ("chat", "error-frame"), ("messages", "error-frame"), ("responses-bridge", "error-frame")], +) +def test_mid_stream_failure_alert_names_the_deployment(alerting: AlertingProxy, entrypoint: str, failure: str) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(chat_upstream(FAILURES[failure], marker)) as wire, alerting.gateway.scenario() as scenario: + model: Final = azure_chat(scenario, wire.url, bridged=entrypoint == "responses-bridge") + caller: Final = keyed_caller(scenario) + seen: Final = ENTRYPOINTS[entrypoint](alerting.gateway, caller.key, model) + assert seen.text == PARTIAL_TEXT and seen.failure, seen + assert failure == "cut" or provider_message(marker) in seen.failure, seen + assert_names_deployment(alert_naming(alerting.alerts, marker), wire.url, model, AZURE_CHAT_BACKEND, caller) + + +def test_native_responses_failed_event_alert_names_the_deployment(alerting: AlertingProxy) -> None: + marker: Final = uuid.uuid4().hex + identity: Final = f"resp_{marker}" + stream: Final = ( + frame(created(identity)), + frame(delta(identity, PARTIAL_TEXT)), + frame(failed(identity, "server_error", provider_message(marker))), + ) + + def respond(request: Request) -> Reply: + if request.method != "POST": + return EMPTY_MODEL_LIST + assert "/responses" in request.target, request.target + return stream_reply(stream) + + with wire_server(respond) as wire, alerting.gateway.scenario() as scenario: + model: Final = scenario.model(model=AZURE_CHAT_BACKEND, api_base=wire.url, api_key=AZURE_KEY, num_retries=0) + caller: Final = keyed_caller(scenario) + seen: Final = responses_outcome(alerting.gateway, caller.key, model) + assert seen.text == PARTIAL_TEXT and provider_message(marker) in seen.failure, seen + assert_names_deployment(alert_naming(alerting.alerts, marker), wire.url, model, AZURE_CHAT_BACKEND, caller) + + +def failing_turn(marker: str) -> Script: + def script( + response_id: str, + ) -> tuple[dict[str, JsonValue], ...]: # mutable-ok: responses_peer's Script returns dicts + opened, text_delta, completed_event = events(response_id, PARTIAL_TEXT) + response: Final = object_value(completed_event["response"]) + return ( + opened, + text_delta, + { + "type": "response.failed", + "response": { + **response, + "status": "failed", + "output": [], + "error": {"code": "server_error", "message": provider_message(marker)}, + }, + }, + ) + + return script + + +def test_responses_websocket_turn_failure_alert_names_the_deployment(alerting: AlertingProxy) -> None: + marker: Final = uuid.uuid4().hex + backend: Final = f"openai/{PROVIDER_MODEL}" + with responses_peer(alerting.cert, failing_turn(marker)) as peer, alerting.gateway.scenario() as scenario: + model: Final = scenario.model(model=backend, api_base=peer.url, num_retries=0) + caller: Final = keyed_caller(scenario) + (received,) = asyncio.run(session(str(alerting.gateway.client.base_url), caller.key, model, (marker,))) + assert [event["type"] for event in received] == [ + "response.created", + "response.output_text.delta", + "response.failed", + ], received + assert_names_deployment(alert_naming(alerting.alerts, marker), peer.url, model, backend, caller) + (row,) = eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (hashlib.sha256(caller.key.encode()).hexdigest(),), + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + assert row["status"] == "failure", row + error_information: Final = object_value(object_value(row["metadata"])["error_information"]) + assert provider_message(marker) in string_value(error_information["error_message"]), error_information + + +@pytest.mark.parametrize( + ("phase", "placement"), + [ + ("pre-stream", "plain"), + ("pre-stream", "code"), + ("pre-stream", "subscription-key"), + ("pre-stream", "userinfo"), + ("mid-stream", "code"), + ], +) +def test_alert_api_base_line_keeps_credentials_out(alerting: AlertingProxy, phase: str, placement: str) -> None: + marker: Final = uuid.uuid4().hex + secret: Final = f"credential{uuid.uuid4().hex}" + pre_stream: Final = phase == "pre-stream" + reply: Final = rejected_before_stream if pre_stream else error_frame_after_text + with wire_server(chat_upstream(reply, marker)) as wire, alerting.gateway.scenario() as scenario: + model: Final = azure_chat(scenario, API_BASES[placement](wire.url, secret)) + caller: Final = keyed_caller(scenario) + seen: Final = chat_outcome(alerting.gateway, caller.key, model) + assert seen.failure and seen.text == ("" if pre_stream else PARTIAL_TEXT), seen + alert: Final = alert_naming(alerting.alerts, model if pre_stream else marker) + assert secret not in alert, alert + assert_names_deployment(alert, f"{wire.url}/gateway", model, AZURE_CHAT_BACKEND, caller) + + +def test_alert_api_base_line_keeps_a_malformed_password_out(alerting: AlertingProxy) -> None: + password_start: Final = f"credential{uuid.uuid4().hex[:8]}" + password_end: Final = uuid.uuid4().hex + with alerting.gateway.scenario() as scenario: + model: Final = azure_chat(scenario, f"http://alice:{password_start}?{password_end}@127.0.0.1:9/gateway") + caller: Final = keyed_caller(scenario) + seen: Final = chat_outcome(alerting.gateway, caller.key, model) + assert seen.failure and not seen.text, seen + line: Final = api_base_line(alert_naming(alerting.alerts, model)) or "" + assert not any(part in line for part in ("alice", password_start, password_end)), line + + +def test_alert_api_base_line_names_a_vertex_partner_model_version(alerting: AlertingProxy) -> None: + with alerting.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"vertex_ai/{VERTEX_PARTNER_BACKEND}", + api_base=None, + api_key=None, + vertex_project=VERTEX_PROJECT, + vertex_location=VERTEX_LOCATION, + vertex_credentials=service_account_json(VERTEX_PROJECT, alerting.gateway.upstream_url.rstrip("/")), + num_retries=0, + ) + caller: Final = keyed_caller(scenario) + seen: Final = chat_outcome(alerting.gateway, caller.key, model) + assert seen.failure and not seen.text, seen + assert any( + f"/publishers/anthropic/models/{VERTEX_PARTNER_BACKEND}:" in request.target + for request in alerting.vertex.drain() + ) + line: Final = api_base_line(alert_naming(alerting.alerts, model)) or "" + assert line.startswith(f"API Base: `{VERTEX_PARTNER_URL}:"), line + + +def test_fallback_mid_stream_failure_alert_names_the_fallback_deployment(alerting: AlertingProxy) -> None: + marker: Final = uuid.uuid4().hex + primary_reply: Final = Reply(status=500, body=b'{"error": {"message": "primary is down", "type": "server_error"}}') + with ( + wire_server(chat_upstream(lambda _: primary_reply, marker)) as primary_wire, + wire_server(chat_upstream(cut_after_text, marker)) as backup_wire, + alerting.gateway.scenario() as scenario, + ): + primary: Final = azure_chat(scenario, primary_wire.url) + backup: Final = azure_chat(scenario, backup_wire.url) + caller: Final = keyed_caller(scenario) + with openai_client(alerting.gateway, caller.key) as client: + items: Final = tuple( + items_or_error( + lambda: client.chat.completions.create( + model=primary, + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + extra_body={**NO_RETRY_NO_CACHE, "fallbacks": [backup]}, + ), + openai.APIError, + ) + ) + seen: Final = outcome(items, chunk_text) + assert seen.text == PARTIAL_TEXT and seen.failure, seen + assert [request.method for request in primary_wire.drain()].count("POST") == 1 + assert [request.method for request in backup_wire.drain()].count("POST") == 1 + alert: Final = alert_naming(alerting.alerts, marker) + assert_names_deployment(alert, backup_wire.url, backup, AZURE_CHAT_BACKEND, caller) + assert primary_wire.url not in alert, alert diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py new file mode 100644 index 00000000000..bb24130aa6c --- /dev/null +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -0,0 +1,433 @@ +import itertools +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_ROW: Final = "anthropic.claude-audit-reload-v1:0" +_GEO: Final = f"au.{_ROW}" +_FINETUNE: Final = "ft:gpt-audit-reload-2025-01-01" +_CLAUDE: Final = "claude-audit-reload-sonnet" +_MINI: Final = "gpt-audit-reload-mini" +_USAGE: Final = re.compile(r"bill (\d+) (\d+) (\d+)") +_CUSTOM: Final = MappingProxyType({"input_cost_per_token": 1.1e-06, "output_cost_per_token": 2.3e-06}) +_COST_MAP_LIMITS: Final = MappingProxyType({"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"}) +_PUBLISHED_SCALE: Final = 2.0 +_CATALOG: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + _ROW: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + }, + _GEO: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + }, + _FINETUNE: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 3.75e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 1.875e-06, + }, + _CLAUDE: { + "litellm_provider": "anthropic", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + }, + _MINI: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + }, + } +) + + +@dataclass(frozen=True, slots=True) +class _Usage: + prompt: int + cached: int + completion: int + + def text(self) -> str: + return f"bill {self.prompt} {self.cached} {self.completion}" + + +@dataclass(frozen=True, slots=True) +class _Rates: + input: float + output: float + cache_read: float + + def cost(self, usage: _Usage) -> float: + return ( + (usage.prompt - usage.cached) * self.input + usage.cached * self.cache_read + usage.completion * self.output + ) + + +@dataclass(frozen=True, slots=True) +class _Deployment: + name: str + model: str + prices: Mapping[str, float] + usage: _Usage + rates: Callable[[float], _Rates | None] + + +@dataclass(frozen=True, slots=True) +class _Charge: + request_id: str + expected: float + + +def _catalog(scale: float) -> dict[str, dict[str, JsonValue]]: + return { + model: {name: value * scale if isinstance(value, float) else value for name, value in row.items()} + for model, row in _CATALOG.items() + } + + +def _rate(model: str, field: str, scale: float) -> float: + rate: Final = _CATALOG[model][field] + assert isinstance(rate, float), rate + return rate * scale + + +def _catalog_rates(model: str) -> Callable[[float], _Rates]: + return lambda scale: _Rates( + input=_rate(model, "input_cost_per_token", scale), + output=_rate(model, "output_cost_per_token", scale), + cache_read=_rate(model, "cache_read_input_token_cost", scale) + if "cache_read_input_token_cost" in _CATALOG[model] + else _rate(model, "input_cost_per_token", scale), + ) + + +def _custom_rates(_scale: float) -> _Rates: + return _Rates( + input=_CUSTOM["input_cost_per_token"], + output=_CUSTOM["output_cost_per_token"], + cache_read=_CUSTOM["input_cost_per_token"], + ) + + +def _unpriced(_scale: float) -> None: + return None + + +def _expected_cost(rates: _Rates | None, usage: _Usage) -> float: + return 0.0 if rates is None else rates.cost(usage) + + +def _expected_listing(rates: _Rates | None) -> tuple[object, object]: + if rates is None: + return None, None + return pytest.approx(rates.input), pytest.approx(rates.output) + + +def _connection(model: str, wire_url: str) -> dict[str, JsonValue]: + if model.startswith("bedrock/"): + return {"api_base": wire_url, "api_key": "synthetic-bedrock-bearer", "aws_region_name": "ap-northeast-1"} + if model.startswith("azure/"): + return {"api_base": wire_url, "api_key": "synthetic-azure-key", "api_version": "2025-01-01-preview"} + return {"api_base": f"{wire_url}/v1", "api_key": "synthetic-openai-key"} + + +def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str) -> Path: + config: Final = directory / "price_data_reload_config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": deployment.name, + "litellm_params": { + "model": deployment.model, + **_connection(deployment.model, wire_url), + **deployment.prices, + }, + } + for deployment in deployments + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _usage_in(request: Request) -> _Usage: + found: Final = _USAGE.search(request.body.decode()) + assert found is not None, request + return _Usage(prompt=int(found.group(1)), cached=int(found.group(2)), completion=int(found.group(3))) + + +def _reply(request: Request) -> Reply: + usage: Final = _usage_in(request) + if "/converse" in request.target: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "priced"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": usage.prompt, + "outputTokens": usage.completion, + "totalTokens": usage.prompt + usage.completion, + }, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + if "/invoke" in request.target: + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [{"type": "text", "text": "priced"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": usage.prompt, "output_tokens": usage.completion}, + } + ).encode() + ) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "scripted", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "priced"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": usage.prompt, + "completion_tokens": usage.completion, + "total_tokens": usage.prompt + usage.completion, + "prompt_tokens_details": {"cached_tokens": usage.cached}, + }, + } + ).encode() + ) + + +def _responder(published: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/model_prices.json": + return Reply(body=json.dumps(_catalog(_PUBLISHED_SCALE if published.is_set() else 1.0)).encode()) + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + return _reply(request) + + return respond + + +def _listed_rates(proxy: Gateway, name: str) -> tuple[JsonValue, JsonValue]: + entries: Final = proxy.get("/model/info")["data"] + assert isinstance(entries, list), entries + info: Final = next( + object_value(object_value(entry)["model_info"]) + for entry in entries + if object_value(entry)["model_name"] == name + ) + return info.get("input_cost_per_token"), info.get("output_cost_per_token") + + +def _billed(client: httpx.Client, key: str, deployment: _Deployment) -> tuple[str, float | None]: + response: Final = client.post( + "/v1/chat/completions", + json={"model": deployment.name, "messages": [{"role": "user", "content": deployment.usage.text()}]}, + headers={"Authorization": f"Bearer {key}"}, + ) + assert response.status_code == 200, response.text + reported: Final = response.headers.get("x-litellm-response-cost") + return string_value(object_value(response.json())["id"]), None if reported is None else float(reported) + + +def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> _Charge: + rates: Final = deployment.rates(scale) + assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name + request_id, charged = _billed(proxy.client, proxy.key, deployment) + expected: Final = _expected_cost(rates, deployment.usage) + assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name + return _Charge(request_id=request_id, expected=expected) + + +def _spend(request_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float)), rows + return float(spend) + + +def _proxy_environment(wire_url: str) -> dict[str, str]: + return {**_COST_MAP_LIMITS, "LITELLM_MODEL_COST_MAP_URL": f"{wire_url}/model_prices.json"} + + +def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments: tuple[_Deployment, ...]) -> None: + published: Final = threading.Event() + with ( + wire_server(_responder(published)) as wire, + owned_proxy( + gateway, + tmp_path, + _proxy_environment(wire.url), + config=_config(tmp_path, deployments, wire.url), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + ) as proxy, + ): + before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in deployments) + published.set() + reload: Final = proxy.request("POST", "/reload/model_cost_map") + assert reload.status_code == 200, reload.text + after: Final = tuple(_charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in deployments) + for charge in (*before, *after): + assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_price_data_reload_reprices_catalog_deployments_and_keeps_custom_prices( + gateway: Gateway, tmp_path: Path +) -> None: + _assert_priced_across_a_reload( + gateway, + tmp_path, + ( + _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)), + _Deployment("mini-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates), + ), + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize( + "deployment", + ( + _Deployment( + "region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW) + ), + _Deployment( + "finetune-unpriced", + f"openai/{_FINETUNE}:org:suffix:y", + {}, + _Usage(2000, 1000, 1000), + _catalog_rates(_FINETUNE), + ), + _Deployment("finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates), + _Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced), + _Deployment( + "region-custom-above-threshold", + f"bedrock/ap-northeast-1.{_ROW}", + _CUSTOM, + _Usage(250_000, 0, 1000), + _custom_rates, + ), + ), + ids=lambda deployment: deployment.name, +) +def test_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload( + gateway: Gateway, tmp_path: Path, deployment: _Deployment +) -> None: + _assert_priced_across_a_reload(gateway, tmp_path, (deployment,)) + + +def _storm_deployment(index: int) -> _Deployment: + prices: Final = MappingProxyType( + {"input_cost_per_token": (index + 1) * 1e-07, "output_cost_per_token": (index + 1) * 3e-07} + ) + return _Deployment( + f"storm-{index}", + f"openai/gpt-audit-reload-storm-{index}", + prices, + _Usage(100, 0, 50), + lambda _scale: _Rates( + input=prices["input_cost_per_token"], + output=prices["output_cost_per_token"], + cache_read=prices["input_cost_per_token"], + ), + ) + + +def _storm_traffic( + base_url: str, key: str, deployments: tuple[_Deployment, ...], offset: int, finished: threading.Event +) -> tuple[tuple[float, float], ...]: + rotation: Final = ( + deployments[(offset + turn) % len(deployments)] for turn, _ in enumerate(iter(finished.is_set, True)) + ) + with httpx.Client(base_url=base_url, timeout=60, trust_env=False) as client: + return tuple( + (_billed(client, key, deployment)[1], _expected_cost(deployment.rates(1.0), deployment.usage)) + for deployment in rotation + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_priced( + gateway: Gateway, tmp_path: Path +) -> None: + storm: Final = tuple(_storm_deployment(index) for index in range(24)) + geo: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)) + finished: Final = threading.Event() + with ( + wire_server(_responder(threading.Event())) as wire, + owned_proxy( + gateway, + tmp_path, + _proxy_environment(wire.url), + config=_config(tmp_path, (geo, *storm), wire.url), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + ) as proxy, + ThreadPoolExecutor(max_workers=8) as pool, + ): + traffic: Final = tuple( + pool.submit(_storm_traffic, str(proxy.client.base_url), proxy.key, storm, offset, finished) + for offset in range(8) + ) + reloads: Final = tuple(proxy.request("POST", "/reload/model_cost_map") for _ in range(40)) + finished.set() + charges: Final = tuple(itertools.chain.from_iterable(calls.result() for calls in traffic)) + assert tuple(reload.status_code for reload in reloads) == (200,) * 40, tuple( + reload.text for reload in reloads if reload.status_code != 200 + ) + assert len(charges) >= 8 + assert tuple(charged for charged, _ in charges) == pytest.approx(tuple(expected for _, expected in charges)) + settled: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (geo, *storm)) + for charge in settled: + assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6) diff --git a/tests/integration/providers/_anthropic_safeguards_beta_support.py b/tests/integration/providers/_anthropic_safeguards_beta_support.py new file mode 100644 index 00000000000..0307ea11724 --- /dev/null +++ b/tests/integration/providers/_anthropic_safeguards_beta_support.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import yaml +from integration._support.anthropic_sse import ( + CONTENT_BLOCK_START, + error_body, + message_start, + sse, + text_delta, +) +from integration._support.client import Gateway, object_value +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue, TypeAdapter + +DANGEROUS_TOOL_USE: Final = "dangerous-tool-use-2026-09-03" +WEB_SEARCH: Final = "web-search-2025-03-05" +CLIENT_BETAS: Final = "interleaved-thinking-2025-05-14,claude-code-20250219,interleaved-thinking-2025-05-14" +MERGED_BETAS: Final = f"claude-code-20250219,{DANGEROUS_TOOL_USE},interleaved-thinking-2025-05-14" +BACKEND: Final = "claude-sonnet-4-6" +ANTHROPIC_DEPLOYMENT: Final = "safeguarded-claude" +VERTEX_DEPLOYMENT: Final = "safeguarded-vertex-claude" +DEEPSEEK_DEPLOYMENT: Final = "safeguarded-deepseek" +ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +DEEPSEEK_KEY: Final = "synthetic-deepseek-key" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-east5" +_VERTEX_MODEL_PATH: Final = f"/v1/projects/{PROJECT}/locations/{LOCATION}/publishers/anthropic/models/{BACKEND}" +VERTEX_TARGET: Final = f"{_VERTEX_MODEL_PATH}:rawPredict" +VERTEX_STREAM_TARGET: Final = f"{_VERTEX_MODEL_PATH}:streamRawPredict?alt=sse" +ANTHROPIC_TARGET: Final = "/v1/messages" +DEEPSEEK_TARGET: Final = "/anthropic/v1/messages" +SCRIPTED_TARGETS: Final = frozenset({ANTHROPIC_TARGET, DEEPSEEK_TARGET, VERTEX_TARGET, VERTEX_STREAM_TARGET}) +REPLY_TEXT: Final = "safeguards beta control" +EXTRA_INPUTS_MESSAGE: Final = "safeguards: Extra inputs are not permitted" +NOT_A_LIST_MESSAGE: Final = "safeguards: Input should be a valid array" +OVERLOADED_MESSAGE: Final = "Overloaded" +SAFEGUARDS: Final[tuple[JsonValue, ...]] = ( + { + "type": "dangerous_tool_use", + "classifier_context": {"v": 1, "permission_mode": "auto"}, + }, +) +SAFEGUARD_RESULTS: Final[tuple[JsonValue, ...]] = ( + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}, +) +BASH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "name": "Bash", + "description": "Run a shell command", + "input_schema": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + } +) +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "web_search_20250305", "name": "web_search", "max_uses": 1} +) +RUST_ROUTE_HEADER: Final = "x-litellm-rust" +_MARKER: Final = re.compile(rb"marker-[0-9a-f]{32}") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = "Started server process [" + + +@dataclass(frozen=True, slots=True) +class Delivery: + target: str + beta: str | None + safeguards: JsonValue + marker: str + authorization: str | None + api_key: str | None + + +def new_marker() -> str: + return f"marker-{uuid.uuid4().hex}" + + +def marker_of(request: Request) -> str: + found: Final = tuple(match.group(0) for match in _MARKER.finditer(request.body)) + return found[-1].decode() if found else "unmarked" + + +def body_of(request: Request) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def json_object(text: str) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(text) + + +def betas_of(request: Request) -> frozenset[str]: + return frozenset(part.strip() for part in request.headers.get("anthropic-beta", "").split(",") if part.strip()) + + +def started_workers(log: Path) -> tuple[int, ...]: + return tuple( + int(line.rsplit("[", 1)[1].strip(" ]\r")) for line in log.read_text().splitlines() if _STARTED_WORKER in line + ) + + +def _streams(request: Request) -> bool: + return request.target == VERTEX_STREAM_TARGET or body_of(request).get("stream") is True + + +def _results(safeguards: JsonValue) -> list[JsonValue] | None: + if safeguards is None: + return None + return list(SAFEGUARD_RESULTS) if safeguards else [] + + +def _message(identity: str, results: list[JsonValue] | None) -> dict[str, JsonValue]: + verdict: Final[dict[str, JsonValue]] = {} if results is None else {"safeguard_results": results} + return { + "id": identity, + "type": "message", + "role": "assistant", + "model": BACKEND, + "content": [{"type": "text", "text": REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + **verdict, + } + + +def _stream(identity: str, results: list[JsonValue] | None) -> tuple[bytes, ...]: + verdict: Final[dict[str, JsonValue]] = {} if results is None else {"safeguard_results": results} + return ( + message_start(identity, BACKEND), + CONTENT_BLOCK_START, + text_delta(REPLY_TEXT), + sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, **verdict}, + "usage": {"output_tokens": 3}, + }, + ), + sse("message_stop", {"type": "message_stop"}), + ) + + +def anthropic_peer(request: Request) -> Reply: + if request.target not in SCRIPTED_TARGETS: + return Reply( + status=404, + body=json.dumps({"error": f"unscripted target {request.target}"}).encode(), + ) + if request.target == DEEPSEEK_TARGET: + return Reply(body=json.dumps(_message(f"msg_{marker_of(request)}", None)).encode()) + body: Final = body_of(request) + safeguards: Final = body.get("safeguards") + if "safeguards" in body and DANGEROUS_TOOL_USE not in betas_of(request): + return Reply(status=400, body=error_body(400, EXTRA_INPUTS_MESSAGE)) + if safeguards is not None and not isinstance(safeguards, list): + return Reply(status=400, body=error_body(400, NOT_A_LIST_MESSAGE)) + identity: Final = f"msg_{marker_of(request)}" + results: Final = _results(safeguards) + if _streams(request): + return Reply(content_type="text/event-stream", chunks=_stream(identity, results)) + return Reply(body=json.dumps(_message(identity, results)).encode()) + + +def outage_peer(outage: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if outage.is_set(): + return Reply(status=529, body=error_body(529, OVERLOADED_MESSAGE)) + return anthropic_peer(request) + + return respond + + +def holding_peer(held: SimpleQueue[str], release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + held.put(marker_of(request)) + assert release.wait(timeout=60), "Held request was never released" + return anthropic_peer(request) + + return respond + + +def deliveries(wire: Wire) -> tuple[Delivery, ...]: + return tuple( + Delivery( + request.target, + request.headers.get("anthropic-beta"), + body_of(request).get("safeguards"), + marker_of(request), + request.headers.get("authorization"), + request.headers.get("x-api-key"), + ) + for request in wire.drain() + ) + + +def one_delivery(wire: Wire, marker: str) -> Delivery: + (sent,) = deliveries(wire) + assert sent.marker == marker, sent + return sent + + +def rust_routed(headers: Mapping[str, str]) -> bool: + return headers.get(RUST_ROUTE_HEADER) == "true" + + +def messages_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": marker}], + "tools": [dict(BASH_TOOL)], + } + + +def safeguarded(body: Mapping[str, JsonValue], safeguards: JsonValue) -> dict[str, JsonValue]: + return {**body, "safeguards": safeguards} + + +def chat_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}]} + + +def responses_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": marker} + + +def owned_config(directory: Path, gateway: Gateway, api_base: str) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + merged: Final = { + **config, + "model_list": [ + { + "model_name": ANTHROPIC_DEPLOYMENT, + "litellm_params": { + "model": f"anthropic/{BACKEND}", + "api_base": api_base, + "api_key": ANTHROPIC_KEY, + }, + }, + { + "model_name": VERTEX_DEPLOYMENT, + "litellm_params": { + "model": f"vertex_ai/{BACKEND}", + "api_base": api_base, + "vertex_project": PROJECT, + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(PROJECT, gateway.upstream_url.rstrip("/")), + }, + }, + { + "model_name": DEEPSEEK_DEPLOYMENT, + "litellm_params": { + "model": "deepseek/deepseek-chat", + "api_base": api_base, + "api_key": DEEPSEEK_KEY, + }, + }, + ], + "router_settings": { + **object_value(config["router_settings"]), + "num_retries": 0, + }, + } + path: Final = directory / f"safeguards-beta-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py index 24063d8052a..37aa96c2467 100644 --- a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py @@ -519,10 +519,12 @@ def test_killing_one_worker_leaves_the_sibling_serving_the_fallback(fb_rig: Rig) workers: Final = eventually(fb_rig.owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) with fb_rig.owned.replacing_worker(min(workers)) as kill: kill() - responses: Final = tuple( - fb_rig.gateway.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill"))) - for _ in range(3) - ) + with httpx.Client(base_url=fb_rig.gateway.client.base_url, timeout=15, trust_env=False) as fresh: + survivor: Final = Gateway(fresh, fb_rig.gateway.key, fb_rig.gateway.upstream_url) + responses: Final = tuple( + survivor.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill"))) + for _ in range(3) + ) for response in responses: assert response.status_code == 200, response.text assert _FALLBACK_TEXT in response.text, response.text diff --git a/tests/integration/providers/test_bedrock_native_structured_output_wire.py b/tests/integration/providers/test_bedrock_native_structured_output_wire.py new file mode 100644 index 00000000000..f389e53d132 --- /dev/null +++ b/tests/integration/providers/test_bedrock_native_structured_output_wire.py @@ -0,0 +1,820 @@ +import json +import re +import uuid +from collections.abc import Iterable, Iterator, Mapping +from itertools import chain +from typing import Final +from urllib.parse import unquote + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, list_value, object_value, string_value +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_HAIKU: Final = "us.anthropic.claude-haiku-5-5" +_FABLE: Final = "us.anthropic.claude-fable-5" +_SONNET: Final = "us.anthropic.claude-sonnet-4-6" +_OPUS: Final = "us.anthropic.claude-opus-5-5" +_OPT_IN: Final = "au.anthropic.claude-opus-4-8" +_DEEPSEEK: Final = "deepseek.v3.2" +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", +} +_NO_CACHE: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}, "num_retries": 0} +_JSON_TOOL: Final = "json_tool_call" +_USER_TOOL: Final = "lookup" +_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" +_ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"} +_THINKING_TYPES: Final = frozenset({"enabled", "adaptive"}) +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_USAGE: Final[dict[str, JsonValue]] = {"inputTokens": 11, "outputTokens": 7, "totalTokens": 18} +_MARKER: Final = re.compile(r"probe ([0-9a-f]{32})") +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_FABLE_REJECTION: Final = ( + "The model returned the following errors: output_config.format: Extra inputs are not permitted" +) +_UNSUPPORTED_KEYWORDS: Final = ( + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "maxItems", + "uniqueItems", + "minProperties", + "maxProperties", + "patternProperties", + "oneOf", +) +_SCHEMA_MAPS: Final = ("properties", "$defs", "definitions") +_SCHEMA_LISTS: Final = ("anyOf", "allOf", "oneOf") +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_LIMITS_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "tier": {"type": "integer", "minimum": 1, "maximum": 5}, + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "owner": {"anyOf": [{"$ref": "#/$defs/Person"}, {"type": "null"}]}, + "backup": {"type": ["object", "null"], "properties": {"name": {"type": "string"}}, "required": ["name"]}, + }, + "required": ["tier", "summary", "tags", "owner", "backup"], + "$defs": {"Person": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}}, +} +_OPEN_COUNTS: Final[dict[str, JsonValue]] = {"type": "object", "additionalProperties": {"type": "integer"}} +_DICT_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"store": {"type": "string"}, "counts": _OPEN_COUNTS}, + "required": ["store", "counts"], + "additionalProperties": False, +} +_NULLABLE_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "backup": {"type": ["object", "null"], "properties": {"name": {"type": "string"}}, "required": ["name"]}, + }, + "required": ["summary", "backup"], + "additionalProperties": False, +} +_CLOSED_NULLABLE_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "backup": { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, + }, + }, + "required": ["summary", "backup"], + "additionalProperties": False, +} +_UNCLEANED_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer", "minimum": 1, "maximum": 5}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_NUMERIC_FORMATS_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "tier": {"type": "integer", "format": "int32"}, + "score": {"type": "number", "format": "double"}, + "summary": {"type": "string"}, + }, + "required": ["tier", "score", "summary"], + "additionalProperties": False, +} +_WIDE_FIELDS: Final = tuple(f"field_{index:02d}" for index in range(65)) +_WIDE_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {name: {"type": "string"} for name in _WIDE_FIELDS}, + "required": list(_WIDE_FIELDS), + "additionalProperties": False, +} +_USER_TOOL_SPEC: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": _USER_TOOL, + "description": "Look up a record", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}, + }, +} + + +def _payload(marker_hex: str) -> dict[str, JsonValue]: + return {"tier": 2, "summary": marker_hex} + + +def _marker() -> tuple[str, str]: + marker_hex: Final = uuid.uuid4().hex + return f"probe {marker_hex}", marker_hex + + +def _response_format(schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "json_schema", "json_schema": {"name": "risk", "schema": dict(schema)}} + + +def _model_id(request: Request) -> str: + return unquote(request.target).split("/model/", 1)[1].rsplit("/", 1)[0] + + +def _is_invoke(request: Request) -> bool: + return unquote(request.target).rsplit("/", 1)[1].startswith("invoke") + + +def _is_stream(request: Request) -> bool: + return unquote(request.target).endswith(("/converse-stream", "/invoke-with-response-stream")) + + +def _thinking_of(body: Mapping[str, JsonValue], invoke: bool) -> JsonValue: + if invoke: + return body.get("thinking") + return object_value(body.get("additionalModelRequestFields") or {}).get("thinking") + + +def _thinking_requested(body: Mapping[str, JsonValue], invoke: bool) -> bool: + thinking: Final = _thinking_of(body, invoke) + return isinstance(thinking, dict) and thinking.get("type") in _THINKING_TYPES + + +def _native_schema(body: Mapping[str, JsonValue], invoke: bool) -> dict[str, JsonValue] | None: + if invoke: + output_config: Final = body.get("output_config") + if not isinstance(output_config, dict) or "format" not in output_config: + return None + return object_value(object_value(object_value(output_config)["format"])["schema"]) + text_format: Final = object_value(body.get("outputConfig") or {}).get("textFormat") + if text_format is None: + return None + json_schema: Final = object_value(object_value(object_value(text_format)["structure"])["jsonSchema"]) + return _JSON.validate_json(string_value(json_schema["schema"])) + + +def _tools(body: Mapping[str, JsonValue], invoke: bool) -> dict[str, dict[str, JsonValue]]: + if invoke: + return { + string_value(object_value(tool)["name"]): object_value(object_value(tool)["input_schema"]) + for tool in list_value(body.get("tools", [])) + } + specs: Final = tuple( + object_value(object_value(tool)["toolSpec"]) + for tool in list_value(object_value(body.get("toolConfig") or {"tools": []})["tools"]) + ) + return {string_value(spec["name"]): object_value(object_value(spec["inputSchema"])["json"]) for spec in specs} + + +def _is_object_type(node: Mapping[str, JsonValue]) -> bool: + node_type: Final = node.get("type") + return node_type == "object" or (isinstance(node_type, list) and "object" in node_type) + + +def _child_schemas(node: Mapping[str, JsonValue]) -> Iterable[JsonValue]: + mapped: Final = (node.get(key) for key in _SCHEMA_MAPS) + listed: Final = (node.get(key) for key in _SCHEMA_LISTS) + single: Final = (node.get(key) for key in ("items", "additionalProperties")) + return chain( + chain.from_iterable(children.values() for children in mapped if isinstance(children, dict)), + chain.from_iterable(variants for variants in listed if isinstance(variants, list)), + (child for child in single if isinstance(child, dict)), + ) + + +def _schema_nodes(schema: JsonValue) -> Iterator[dict[str, JsonValue]]: + if not isinstance(schema, dict): + return + yield schema + yield from chain.from_iterable(_schema_nodes(child) for child in _child_schemas(schema)) + + +def _node_fault(node: Mapping[str, JsonValue]) -> str | None: + node_type: Final = node.get("type") + unsupported: Final = tuple(key for key in _UNSUPPORTED_KEYWORDS if key in node) + if unsupported: + return f"For '{node_type}' type, property '{unsupported[0]}' is not supported" + min_items: Final = node.get("minItems") + if isinstance(min_items, int) and min_items > 1: + return "For 'array' type, property 'minItems' is not supported" + if _is_object_type(node) and node.get("additionalProperties") is not False: + return "For 'object' type, 'additionalProperties' must be explicitly set to false" + return None + + +def _bedrock_schema_fault(schema: JsonValue) -> str | None: + faults: Final = tuple(fault for node in _schema_nodes(schema) if (fault := _node_fault(node)) is not None) + if not faults: + return None + return f"The model returned the following errors: output_config.format.schema: {faults[0]}" + + +def _rejection(model_id: str, native: dict[str, JsonValue] | None) -> str | None: + if native is None: + return None + if model_id == _FABLE: + return _FABLE_REJECTION + if "anthropic." not in model_id: + return None + return _bedrock_schema_fault(native) + + +def _reasoning_block(marker_hex: str, invoke: bool) -> dict[str, JsonValue]: + if invoke: + return {"type": "thinking", "thinking": f"reasoning {marker_hex}", "signature": _SIGNATURE} + return {"reasoningContent": {"reasoningText": {"text": f"reasoning {marker_hex}", "signature": _SIGNATURE}}} + + +def _answer_block(name: str | None, marker_hex: str, invoke: bool) -> dict[str, JsonValue]: + tool_input: Final[dict[str, JsonValue]] = {"q": marker_hex} if name == _USER_TOOL else _payload(marker_hex) + if name is None and invoke: + return {"type": "text", "text": json.dumps(_payload(marker_hex))} + if name is None: + return {"text": json.dumps(_payload(marker_hex))} + if invoke: + return {"type": "tool_use", "id": f"toolu_{marker_hex}", "name": name, "input": tool_input} + return {"toolUse": {"toolUseId": f"tooluse_{marker_hex}", "name": name, "input": tool_input}} + + +def _converse_frames(blocks: tuple[dict[str, JsonValue], ...], stop_reason: str) -> bytes: + def frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload).encode(), + ) + + def block_frames(index: int, block: dict[str, JsonValue]) -> bytes: + if "reasoningContent" in block: + reasoning: Final = object_value(object_value(block["reasoningContent"])["reasoningText"]) + return b"".join( + ( + frame( + "contentBlockDelta", + {"delta": {"reasoningContent": {"text": reasoning["text"]}}, "contentBlockIndex": index}, + ), + frame( + "contentBlockDelta", + { + "delta": {"reasoningContent": {"signature": reasoning["signature"]}}, + "contentBlockIndex": index, + }, + ), + frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ) + if "toolUse" in block: + tool_use: Final = object_value(block["toolUse"]) + return b"".join( + ( + frame( + "contentBlockStart", + { + "start": {"toolUse": {"toolUseId": tool_use["toolUseId"], "name": tool_use["name"]}}, + "contentBlockIndex": index, + }, + ), + frame( + "contentBlockDelta", + {"delta": {"toolUse": {"input": json.dumps(tool_use["input"])}}, "contentBlockIndex": index}, + ), + frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ) + return b"".join( + ( + frame("contentBlockDelta", {"delta": {"text": block["text"]}, "contentBlockIndex": index}), + frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ) + + return b"".join( + ( + frame("messageStart", {"role": "assistant"}), + *(block_frames(index, block) for index, block in enumerate(blocks)), + frame("messageStop", {"stopReason": stop_reason}), + frame("metadata", {"usage": _USAGE, "metrics": {"latencyMs": 1}}), + ) + ) + + +def _called_tool(tools: Iterable[str], native: bool) -> str | None: + names: Final = frozenset(tools) + if _USER_TOOL in names: + return _USER_TOOL + if native or _JSON_TOOL not in names: + return None + return _JSON_TOOL + + +def _error_reply(message: str) -> Reply: + return Reply(status=400, body=json.dumps({"message": message}).encode()) + + +def _bedrock(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + invoke: Final = _is_invoke(request) + found: Final = _MARKER.search(json.dumps(body)) + assert found is not None, body + marker_hex: Final = found.group(1) + native: Final = _native_schema(body, invoke) + rejection: Final = _rejection(_model_id(request), native) + if rejection is not None: + return _error_reply(rejection) + called: Final = _called_tool(_tools(body, invoke).keys(), native is not None) + reasoning: Final = (_reasoning_block(marker_hex, invoke),) if _thinking_requested(body, invoke) else () + blocks: Final = (*reasoning, _answer_block(called, marker_hex, invoke)) + stop_reason: Final = "tool_use" if called is not None else "end_turn" + if invoke: + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{marker_hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": list(blocks), + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 7}, + } + return Reply(body=json.dumps(message).encode()) + if _is_stream(request): + return Reply(body=_converse_frames(blocks, stop_reason), content_type=_EVENT_STREAM) + converse: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": list(blocks)}}, + "stopReason": stop_reason, + "usage": _USAGE, + "metrics": {"latencyMs": 1}, + } + return Reply(body=json.dumps(converse).encode()) + + +def _deployment( + scenario: Scenario, wire: Wire, route: str, backend: str, model_info: Mapping[str, JsonValue] | None = None +) -> str: + return scenario.model(model=f"bedrock/{route}/{backend}", api_base=wire.url, model_info=model_info, **_AWS) + + +def _chat_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], "max_tokens": 2048, **_NO_CACHE, **extra} + + +def _messages_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "max_tokens": 2048, "messages": [{"role": "user", "content": marker}], **_NO_CACHE, **extra} + + +def _one_call(wire: Wire, marker_hex: str) -> tuple[Request, dict[str, JsonValue]]: + (request,) = wire.drain() + body: Final = _JSON.validate_json(request.body) + assert marker_hex in json.dumps(body), body + return request, body + + +def _chat_message(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + choices: Final = list_value(_JSON.validate_json(response.content)["choices"]) + return object_value(object_value(choices[0])["message"]) + + +def _stream_message(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + deltas: Final = tuple( + object_value(object_value(list_value(_JSON.validate_json(line[6:])["choices"])[0])["delta"]) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + return { + "content": "".join(string_value(delta["content"]) for delta in deltas if delta.get("content")), + "reasoning_content": "".join( + string_value(delta["reasoning_content"]) for delta in deltas if delta.get("reasoning_content") + ), + } + + +def _content_json(message: Mapping[str, JsonValue]) -> JsonValue: + return _JSON_VALUE.validate_json(string_value(message["content"])) + + +def _responses_items(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + assert response.status_code == 200, response.text + return tuple(object_value(item) for item in list_value(_JSON.validate_json(response.content)["output"])) + + +def _responses_text(items: tuple[dict[str, JsonValue], ...]) -> JsonValue: + (message,) = tuple(item for item in items if item.get("type") == "message") + (text,) = tuple( + object_value(part) for part in list_value(message["content"]) if object_value(part).get("type") == "output_text" + ) + return _JSON_VALUE.validate_json(string_value(text["text"])) + + +def _messages_blocks(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + assert response.status_code == 200, response.text + return tuple(object_value(block) for block in list_value(_JSON.validate_json(response.content)["content"])) + + +def _messages_text(blocks: tuple[dict[str, JsonValue], ...]) -> JsonValue: + (text,) = tuple(block for block in blocks if block.get("type") == "text") + return _JSON_VALUE.validate_json(string_value(text["text"])) + + +def _refs(schema: JsonValue) -> frozenset[str]: + return frozenset(string_value(node["$ref"]) for node in _schema_nodes(schema) if "$ref" in node) + + +def _resolves(schema: Mapping[str, JsonValue], ref: str) -> bool: + prefix: Final = "#/$defs/" + return ref.startswith(prefix) and ref.removeprefix(prefix) in object_value(schema.get("$defs") or {}) + + +def _assert_native_without_tools(body: Mapping[str, JsonValue], invoke: bool, schema: Mapping[str, JsonValue]) -> None: + assert _native_schema(body, invoke) == schema, body + assert _JSON_TOOL not in _tools(body, invoke), body + + +@pytest.mark.parametrize( + ("route", "stream"), + [("converse", False), ("converse", True), ("invoke", False)], + ids=["converse", "converse_stream", "invoke"], +) +def test_thinking_request_on_a_bedrock_native_model_gets_native_output_and_keeps_reasoning( + gateway: Gateway, route: str, stream: bool +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, stream=stream, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)), + ) + message: Final = _stream_message(response) if stream else _chat_message(response) + assert _content_json(message) == _payload(marker_hex), response.text + assert message["reasoning_content"] == f"reasoning {marker_hex}", response.text + request, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert _is_stream(request) is stream, request.target + _assert_native_without_tools(body, invoke, _SCHEMA) + assert _thinking_requested(body, invoke), body + + +def test_responses_reasoning_effort_on_a_bedrock_native_model_gets_native_output_and_keeps_reasoning( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "input": marker, + "max_output_tokens": 2048, + "reasoning": {"effort": "low"}, + "text": {"format": {"type": "json_schema", "name": "risk", "schema": _SCHEMA}}, + **_NO_CACHE, + }, + ) + items: Final = _responses_items(response) + assert _responses_text(items) == _payload(marker_hex), response.text + (reasoning,) = tuple(item for item in items if item.get("type") == "reasoning") + assert f"reasoning {marker_hex}" in json.dumps(reasoning), response.text + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, False, _SCHEMA) + assert _thinking_requested(body, False), body + + +def test_messages_thinking_on_a_bedrock_converse_native_model_gets_native_output_and_keeps_reasoning( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, thinking=_ADAPTIVE, output_format={"type": "json_schema", "schema": _SCHEMA}), + ) + blocks: Final = _messages_blocks(response) + assert _messages_text(blocks) == _payload(marker_hex), response.text + (thinking,) = tuple(block for block in blocks if block.get("type") == "thinking") + assert thinking["thinking"] == f"reasoning {marker_hex}", response.text + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, False, _SCHEMA) + assert _thinking_requested(body, False), body + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_bedrock_native_model_without_thinking_keeps_the_forced_json_tool(gateway: Gateway, route: str) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _HAIKU) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) + ) + message: Final = _chat_message(response) + assert _content_json(message) == _payload(marker_hex), response.text + assert not message.get("tool_calls"), response.text + _, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert _native_schema(body, invoke) is None, body + assert _tools(body, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body + + +@pytest.mark.parametrize( + ("backend", "model_info"), + [(_OPUS, None), (_OPT_IN, {"supports_bedrock_native_structured_output": True})], + ids=["neither_key", "deployment_opt_in"], +) +def test_thinking_request_follows_the_bedrock_native_key_the_model_resolves_to( + gateway: Gateway, backend: str, model_info: Mapping[str, JsonValue] | None +) -> None: + marker, marker_hex = _marker() + native: Final = model_info is not None + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", backend, model_info) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)), + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + assert (_native_schema(body, False) == _SCHEMA) is native, body + assert (_JSON_TOOL in _tools(body, False)) is not native, body + + +def _opt_in_answer(gateway: Gateway, entry: str, model: str, marker: str) -> JsonValue: + output_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": _SCHEMA} + match entry: + case "converse_chat" | "invoke_chat": + chat: Final = _chat_body(model, marker, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", chat))) + case "converse_responses": + responses: Final[dict[str, JsonValue]] = { + "model": model, + "input": marker, + "max_output_tokens": 2048, + "reasoning": {"effort": "low"}, + "text": {"format": {**output_format, "name": "risk"}}, + **_NO_CACHE, + } + return _responses_text(_responses_items(gateway.request("POST", "/v1/responses", responses))) + case "converse_messages": + adapted: Final = _messages_body(model, marker, thinking=_ADAPTIVE, output_format=output_format) + return _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", adapted))) + case _: + invoked: Final = _messages_body(model, marker, output_config={"format": output_format}) + return _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", invoked))) + + +def _sent_native_schema(gateway: Gateway, wire: Wire, entry: str, model: str) -> dict[str, JsonValue] | None: + marker, marker_hex = _marker() + assert _opt_in_answer(gateway, entry, model, marker) == _payload(marker_hex) + _, body = _one_call(wire, marker_hex) + invoke: Final = entry.startswith("invoke") + native: Final = _native_schema(body, invoke) + assert native is None or _JSON_TOOL not in _tools(body, invoke), body + return native + + +@pytest.mark.parametrize( + "entry", ["converse_chat", "converse_responses", "converse_messages", "invoke_chat", "invoke_messages"] +) +def test_a_deployment_native_output_opt_in_stays_on_the_deployment_that_sets_it(gateway: Gateway, entry: str) -> None: + route: Final = entry.split("_", 1)[0] + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + opted_in: Final = _deployment( + scenario, wire, route, _OPT_IN, {"supports_bedrock_native_structured_output": True} + ) + sibling: Final = _deployment(scenario, wire, route, _OPT_IN) + assert _sent_native_schema(gateway, wire, entry, sibling) is None + assert _sent_native_schema(gateway, wire, entry, opted_in) == _SCHEMA + + +def test_fable_5_json_schema_takes_the_json_tool_instead_of_the_native_output_bedrock_rejects( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _FABLE) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + assert "outputConfig" not in body, body + assert _tools(body, False)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body + + +def _limits_answer(gateway: Gateway, entry: str, model: str, marker: str) -> JsonValue: + output_config: Final[dict[str, JsonValue]] = {"format": {"type": "json_schema", "schema": _LIMITS_SCHEMA}} + match entry: + case "converse_chat" | "invoke_chat": + body: Final = _chat_body(model, marker, response_format=_response_format(_LIMITS_SCHEMA)) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", body))) + case "converse_output_config": + chat: Final = _chat_body(model, marker, output_config=output_config) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", chat))) + case _: + messages: Final = _messages_body(model, marker, output_config=output_config) + return _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", messages))) + + +@pytest.mark.parametrize("entry", ["converse_chat", "converse_output_config", "invoke_chat", "invoke_messages"]) +def test_native_schema_reaches_bedrock_with_limits_in_descriptions_and_every_object_closed( + gateway: Gateway, entry: str +) -> None: + marker, marker_hex = _marker() + route: Final = entry.split("_", 1)[0] + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _SONNET) + assert _limits_answer(gateway, entry, model, marker) == _payload(marker_hex) + _, body = _one_call(wire, marker_hex) + schema: Final = _native_schema(body, route == "invoke") + assert schema is not None, body + assert _bedrock_schema_fault(schema) is None, schema + properties: Final = object_value(schema["properties"]) + assert object_value(properties["tags"])["minItems"] == 1, schema + assert {"1", "5"} <= set(re.findall(r"\d+", string_value(object_value(properties["tier"])["description"]))) + assert all(_resolves(schema, ref) for ref in _refs(schema)), schema + assert _JSON_TOOL not in _tools(body, route == "invoke"), body + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_open_object_schema_takes_the_json_tool_on_a_native_model_and_keeps_the_dict_open( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _SONNET) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_DICT_SCHEMA)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert _native_schema(body, invoke) is None, body + tool_schema: Final = _tools(body, invoke)[_JSON_TOOL] + assert object_value(tool_schema["properties"])["counts"] == _OPEN_COUNTS, tool_schema + + +def test_thinking_with_a_caller_tool_on_a_bedrock_native_model_sends_only_the_callers_tool(gateway: Gateway) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + marker, + thinking=_ADAPTIVE, + tools=[_USER_TOOL_SPEC], + response_format=_response_format(_SCHEMA), + ), + ) + (call,) = list_value(_chat_message(response)["tool_calls"]) + function: Final = object_value(object_value(call)["function"]) + assert function["name"] == _USER_TOOL, response.text + assert json.loads(string_value(function["arguments"])) == {"q": marker_hex}, response.text + _, body = _one_call(wire, marker_hex) + assert _tools(body, False).keys() == {_USER_TOOL}, body + assert _native_schema(body, False) == _SCHEMA, body + assert _thinking_requested(body, False), body + + +@pytest.mark.parametrize( + ("schema", "expected"), + [(_UNCLEANED_SCHEMA, _UNCLEANED_SCHEMA), (_NULLABLE_SCHEMA, _CLOSED_NULLABLE_SCHEMA)], + ids=["limits_kept", "nullable_object_closed"], +) +def test_non_anthropic_native_model_gets_the_callers_schema_with_objects_closed( + gateway: Gateway, schema: dict[str, JsonValue], expected: dict[str, JsonValue] +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _DEEPSEEK) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(schema)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + assert _native_schema(body, False) == expected, body + + +@pytest.mark.parametrize( + ("schema", "native"), + [ + (_SCHEMA, True), + (_DICT_SCHEMA, False), + ], + ids=["closed_schema_native", "open_object_inlined"], +) +def test_messages_on_bedrock_invoke_follows_the_bedrock_native_key_without_thinking( + gateway: Gateway, schema: dict[str, JsonValue], native: bool +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "invoke", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_config={"format": {"type": "json_schema", "schema": schema}}), + ) + assert _messages_text(_messages_blocks(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + sent: Final = _native_schema(body, True) + assert (sent is not None) is native, body + prompt: Final = json.dumps(body["messages"]) + assert all((name in prompt) is not native for name in object_value(schema["properties"])), body + + +@pytest.mark.parametrize( + "schema", + [ + {"type": "object", "properties": ["tier", "summary"]}, + {"type": "object", "properties": {"tier": {"$ref": 5}}, "additionalProperties": False}, + {"type": "object", "$defs": [1, 2], "properties": {"tier": {"$ref": "#/$defs/Tier"}}}, + {"type": ["object", "null"]}, + {"type": "object", "properties": {"tier": {"type": "integer"}}, "required": "tier"}, + ], + ids=["properties_list", "ref_not_string", "defs_list", "type_list_without_properties", "required_string"], +) +def test_malformed_schema_on_a_native_model_reaches_bedrock_once_without_a_proxy_error( + gateway: Gateway, schema: dict[str, JsonValue] +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _SONNET) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(schema)) + ) + assert response.status_code < 500, response.text + _, body = _one_call(wire, marker_hex) + native: Final = _native_schema(body, False) + fault: Final = None if native is None else _bedrock_schema_fault(native) + assert response.status_code == (200 if fault is None else 400), response.text + assert fault is None or fault.split(": ", 1)[1] in response.text, response.text + assert (native is None) is (_JSON_TOOL in _tools(body, False)), body + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize( + ("route", "backend", "thinking", "schema", "native"), + [ + ("converse", _SONNET, False, _NUMERIC_FORMATS_SCHEMA, True), + ("invoke", _HAIKU, True, _NUMERIC_FORMATS_SCHEMA, True), + ("converse", _SONNET, False, _WIDE_SCHEMA, True), + ("converse", _HAIKU, True, _WIDE_SCHEMA, False), + ], + ids=["numeric_formats_old_key", "numeric_formats_new_key_invoke", "wide_old_key", "wide_new_key"], +) +def test_numeric_formats_and_wide_schemas_follow_the_bedrock_native_key( + gateway: Gateway, route: str, backend: str, thinking: bool, schema: dict[str, JsonValue], native: bool +) -> None: + marker, marker_hex = _marker() + extra: Final[dict[str, JsonValue]] = {"thinking": _ADAPTIVE} if thinking else {} + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, backend) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, response_format=_response_format(schema), **extra), + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert (_native_schema(body, invoke) == schema) is native, body + tools: Final = _tools(body, invoke) + assert (_JSON_TOOL in tools) is not native, body + assert native or tools[_JSON_TOOL]["properties"] == schema["properties"], body diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py index ccef839a163..1c8850b9bfa 100644 --- a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -8,6 +8,11 @@ dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream) sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult documents the proxy forwarded to the peer. +The fallback cells cover a websocket that outlives one provider attempt. A refused stream keeps the +client's opening message for the fallback, which gets it once whether it dials Bedrock again or +an OpenAI realtime deployment, and once Bedrock accepted a stream no later attempt may open a second +provider session on that websocket. + The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE`` pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it. """ @@ -21,12 +26,19 @@ import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager from pathlib import Path +from types import MappingProxyType from typing import Final +import httpx import pytest import websockets +import yaml from integration._support.bedrock_realtime_peer import ( + BedrockEvent, BedrockRealtimePeer, + OpeningEnd, + Refusal, + StreamFailure, Turn, assistant_text, bedrock_realtime_peer, @@ -43,12 +55,15 @@ from integration._support.client import ( object_value, string_value, ) -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import write_self_signed_cert +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import RealtimeResponse from pydantic import JsonValue from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} @@ -60,6 +75,30 @@ WEATHER_FUNCTION: Final = { NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}' +QUESTION: Final = "What is the weather in New York right now?" +ANSWER: Final = "It is 72 and sunny" +INSTRUCTIONS: Final = "Answer every weather question in one short sentence." +OPENAI_VOICE_MODEL: Final = "gpt-realtime-2.1" +BEDROCK_REPLAY_GROUPS: Final = MappingProxyType( + { + "system_prompt": ("sonic-refused-session-update-first", "sonic-replay-session-update-first"), + "user_input": ("sonic-refused-text-first", "sonic-replay-text-first"), + } +) +OPENAI_REPLAY_GROUPS: Final = ("sonic-refused-openai-fallback", "openai-replay-target") +OPENAI_COMMITTED_GROUPS: Final = ("sonic-committed-openai-fallback", "openai-committed-target") +FALLBACK_GROUPS: Final = (*BEDROCK_REPLAY_GROUPS.values(), OPENAI_REPLAY_GROUPS, OPENAI_COMMITTED_GROUPS) +SCRIPTED_DONE: Final[dict[str, JsonValue]] = { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": { + "id": "resp_$UNIQUE_ID", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}, + }, +} @pytest.fixture(scope="module") @@ -67,15 +106,35 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert")) +def _fallback_config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python( + yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text()) + ) + config: Final = { + **stock, + "router_settings": { + **object_value(stock["router_settings"]), + "fallbacks": [{source: [target]} for source, target in FALLBACK_GROUPS], + }, + } + path: Final = directory / "bedrock-sonic-fallbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + @pytest.fixture(scope="module") def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("bedrock-sonic") with gateway_from_environment() as base: - with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_fallback_config(directory)) as proxy: yield proxy def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: - name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}" + return _bedrock_deployment(gateway, scenario, peer, f"bedrock-sonic-{uuid.uuid4().hex[:12]}") + + +def _bedrock_deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer, name: str) -> str: created: Final = gateway.post( "/model/new", { @@ -94,6 +153,44 @@ def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) return name +def _openai_deployment(gateway: Gateway, scenario: Scenario, name: str) -> str: + scenario_id: Final = f"bedrock-sonic-fallback-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario( + scenario_id, RealtimeResponse(content_type="application/x-realtime", events=(SCRIPTED_DONE,)) + ) + scenario.cleanups.callback(delete_scenario, handle) + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": f"openai/{OPENAI_VOICE_MODEL}", + "api_key": scenario_id, + "api_base": f"{gateway.upstream_url}/v1", + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return scenario_id + + +def _upstream_requests(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content) + requests: Final = observations["requests"] + assert isinstance(requests, list), observations + return tuple( + request + for request in map(object_value, requests) + if request["authorization"] == f"Bearer {scenario_id}" or request["api_key"] == scenario_id + ) + + +def _upstream_frames(requests: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(object_value(request["body"]) for request in requests if request["method"] == "WEBSOCKET_FRAME") + + def _session_update(tool: Mapping[str, JsonValue]) -> str: return json.dumps( { @@ -113,12 +210,16 @@ def _user_text() -> str: "item": { "type": "message", "role": "user", - "content": [{"type": "input_text", "text": "What is the weather in New York right now?"}], + "content": [{"type": "input_text", "text": QUESTION}], }, } ) +def _instructions_update() -> str: + return json.dumps({"type": "session.update", "session": {"instructions": INSTRUCTIONS}}) + + _RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) @@ -207,10 +308,13 @@ def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple assert _tail_types(events) == ( "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.done", ), events + (delta,) = tuple(event for event in events if event.get("type") == "response.function_call_arguments.delta") + assert string_value(delta["delta"]) == '{"location":"New York"}' (call,) = _function_calls(events) assert string_value(call["call_id"]) == "tooluse_A" assert string_value(call["name"]) == "get_weather" @@ -376,3 +480,160 @@ def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], con events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) (call,) = _function_calls(events) assert string_value(call["arguments"]) == (content if content is not None else "{}"), events + + +def _of_type(events: tuple[Event, ...], event_type: str) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == event_type) + + +def _bedrock_fields(events: tuple[BedrockEvent, ...], name: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(object_value(event["event"]) for event in events) + return tuple(object_value(body[name]) for body in bodies if name in body) + + +def _text_inputs(events: tuple[BedrockEvent, ...], role: str) -> tuple[str, ...]: + roles: Final = MappingProxyType( + {string_value(start["contentName"]): start.get("role") for start in _bedrock_fields(events, "contentStart")} + ) + return tuple( + string_value(text["content"]) + for text in _bedrock_fields(events, "textInput") + if roles.get(string_value(text["contentName"])) == role + ) + + +def _answer_turn() -> Turn: + return Turn("user_input", assistant_text(ANSWER) + (usage(30, 12),)) + + +def _answer_text(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in _of_type(events, "response.text.delta")) + + +async def _opened_with(socket: ClientConnection, opening: tuple[str, ...], stop: frozenset[str]) -> tuple[Event, ...]: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + for message in opening: + await socket.send(message) + return (created,) + await _events_until(socket, lambda event, _events: event.get("type") in stop) + + +async def _session( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...], stop: frozenset[str] +) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + return await _opened_with(socket, opening, stop) + + +async def _session_until_closed( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...] +) -> tuple[tuple[Event, ...], int | None]: + async with _realtime(gateway, model, key) as socket: + events: Final = await _opened_with(socket, opening, frozenset(("error", "session.created"))) + if events[-1].get("type") != "error": + return events, None + try: + await asyncio.wait_for(socket.recv(), 30) + except ConnectionClosed: + return events, socket.close_code + return events, None + + +@pytest.mark.parametrize("after", ["system_prompt", "user_input"], ids=["session-update-first", "text-first"]) +def test_refused_opening_is_replayed_on_the_bedrock_fallback( + candidate: Gateway, cert: tuple[Path, Path], after: OpeningEnd +) -> None: + source, target = BEDROCK_REPLAY_GROUPS[after] + update_first: Final = after == "system_prompt" + opening: Final = ((_instructions_update(),) if update_first else ()) + (_user_text(), _RESPONSE_CREATE) + opening_role, opening_text = ("SYSTEM", INSTRUCTIONS) if update_first else ("USER", QUESTION) + with ( + bedrock_realtime_peer((), cert, refusal=Refusal(after)) as refusing, + bedrock_realtime_peer((_answer_turn(),), cert) as healthy, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + _bedrock_deployment(candidate, scenario, healthy, target) + events: Final = asyncio.run( + _session( + candidate, source, scenario.key(), opening, frozenset(("response.done", "error", "session.created")) + ) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_of_type(events, "session.updated")) == int(update_first), events + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + (replayed,) = healthy.streams() + assert _text_inputs(replayed, "USER") == (QUESTION,), replayed + assert _text_inputs(replayed, "SYSTEM").count(INSTRUCTIONS) == int(update_first), replayed + refused: Final = refusing.streams() + assert refused and all(opening_text in _text_inputs(stream, opening_role) for stream in refused), refused + + +def test_refused_text_opening_is_replayed_on_the_openai_fallback(candidate: Gateway, cert: tuple[Path, Path]) -> None: + source, target = OPENAI_REPLAY_GROUPS + with ( + bedrock_realtime_peer((), cert, refusal=Refusal("user_input")) as refusing, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events: Final = asyncio.run( + _session( + candidate, + source, + scenario.key(), + (_user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "error")), + ) + ) + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + requests: Final = _upstream_requests(candidate, scenario_id) + assert len(tuple(request for request in requests if request["method"] == "WEBSOCKET")) == 1, requests + assert _upstream_frames(requests) == (json.loads(_user_text()), json.loads(_RESPONSE_CREATE)), requests + + +def test_committed_session_failure_closes_the_socket_without_another_provider_stream( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + source, target = OPENAI_COMMITTED_GROUPS + with ( + bedrock_realtime_peer((_answer_turn(),), cert, failure=StreamFailure()) as committed, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, committed, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events, close_code = asyncio.run( + _session_until_closed(candidate, source, scenario.key(), (_user_text(), _RESPONSE_CREATE)) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + assert string_value(events[-1]["type"]) == "error", events + assert close_code == 1011, events + assert _upstream_requests(candidate, scenario_id) == () + assert len(committed.streams()) == 1 + + +def test_voice_object_is_refused_and_never_reaches_bedrock(candidate: Gateway, cert: tuple[Path, Path]) -> None: + voice_update: Final = json.dumps({"type": "session.update", "session": {"voice": {"id": "voice_123"}}}) + with bedrock_realtime_peer((_answer_turn(),), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + events: Final = asyncio.run( + _session( + candidate, + model, + scenario.key(), + (voice_update, _user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "session.created")), + ) + ) + errors: Final = tuple(object_value(event["error"]) for event in _of_type(events, "error")) + assert tuple((error.get("code"), error.get("param")) for error in errors) == ( + ("invalid_type", "session.voice"), + ), events + assert _of_type(events, "session.updated") == (), events + assert len(_of_type(events, "session.created")) == 1, events + assert _answer_text(events) == ANSWER, events + assert "voice_123" not in json.dumps(peer.received()) diff --git a/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py new file mode 100644 index 00000000000..5fd12405ab2 --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py @@ -0,0 +1,316 @@ +import base64 +import json +import uuid +from urllib.parse import unquote +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_METADATA_HEADER: Final = "x-amzn-bedrock-request-metadata" +_HEADER_MAX_CHARS: Final = 8500 +_LONG_SPEND_LOGS_METADATA: Final = tuple((f"{index:02d}".ljust(256, "k"), "\n" * 256) for index in range(12)) +_CONVERSE_REPLY: Final[Mapping[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}, + "metrics": {"latencyMs": 3}, +} +_INVOKE_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} +_CONVERSE_STREAM_FRAMES: Final = ( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"), + _aws_event_frame("metadata", {"usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}}, "sc", "u"), +) + + +def _bedrock_peer(request: Request) -> Reply: + if request.target.endswith("/converse-stream"): + return Reply(chunks=_CONVERSE_STREAM_FRAMES, content_type=_EVENT_STREAM) + if request.target.endswith("/invoke-with-response-stream"): + chunk: Final = {"bytes": base64.b64encode(json.dumps(_INVOKE_REPLY).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + if request.target.endswith("/converse"): + return Reply(body=json.dumps(_CONVERSE_REPLY).encode()) + return Reply(body=json.dumps(_INVOKE_REPLY).encode()) + + +def _messages(marker: str) -> Sequence[JsonValue]: + return [{"role": "user", "content": [{"text": marker}]}] + + +def _invoke_body(marker: str) -> Mapping[str, JsonValue]: + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + } + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _header_value(pairs: tuple[tuple[str, str], ...]) -> str: + return json.dumps(dict(pairs), separators=(",", ":")) + + +def _send_with_caller_metadata( + runtime: object, action: str, model: str, marker: str, caller_metadata: Mapping[str, JsonValue] +) -> None: + if action == "invoke": + getattr(runtime, "invoke_model")( + modelId=model, body=json.dumps({**_invoke_body(marker), "requestMetadata": caller_metadata}) + ) + return + getattr(runtime, action.replace("-", "_"))( + modelId=model, messages=_messages(marker), requestMetadata=caller_metadata + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + strict_model: str + dropping_model: str + chat_model: str + key: str + identity: Mapping[str, str] + long_key: str + long_identity: tuple[tuple[str, str], ...] + + def runtime(self, key: str, monkeypatch: pytest.MonkeyPatch): # noqa: ANN201 # boto3 clients are untyped + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key) + return boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config( + read_timeout=case_timeout_seconds(), + retries={"total_max_attempts": 1}, + parameter_validation=False, + ), + ) + + def received(self, marker: str) -> tuple[Request, ...]: + return tuple(request for request in self.wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + strict_model: Final = f"bedrock-metadata-{suffix}" + dropping_model: Final = f"bedrock-metadata-dropping-{suffix}" + chat_model: Final = f"bedrock-metadata-chat-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-request-metadata") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_bedrock_peer) as wire: + deployment: Final = { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + } + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + {"model_name": strict_model, "litellm_params": deployment}, + {"model_name": dropping_model, "litellm_params": {**deployment, "drop_params": True}}, + { + "model_name": chat_model, + "litellm_params": {**deployment, "model": f"bedrock/converse/{_MODEL_ID}", "drop_params": True}, + }, + ], + "litellm_settings": { + **base_config.get("litellm_settings", {}), + "drop_params": False, + "bedrock_request_metadata_fields": [ + "user_api_key_alias", + "user_api_key_team_alias", + "spend_logs_metadata", + ], + }, + } + config_path: Final = directory / "bedrock-request-metadata.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + team_alias: Final = f"platform-{suffix}" + team_id: Final = scenario.team(team_alias=team_alias) + key_alias: Final = f"prod-key-{suffix}" + long_key_alias: Final = f"long-key-{suffix}" + yield Rig( + proxy=proxy, + wire=wire, + strict_model=strict_model, + dropping_model=dropping_model, + chat_model=chat_model, + key=scenario.key( + key_alias=key_alias, + team_id=team_id, + metadata={"spend_logs_metadata": {"agent": "triage-bot"}}, + ), + identity={ + "user_api_key_alias": key_alias, + "user_api_key_team_alias": team_alias, + "agent": "triage-bot", + }, + long_key=scenario.key( + key_alias=long_key_alias, + team_id=team_id, + metadata={"spend_logs_metadata": dict(_LONG_SPEND_LOGS_METADATA)}, + ), + long_identity=( + ("user_api_key_alias", long_key_alias), + ("user_api_key_team_alias", team_alias), + *_LONG_SPEND_LOGS_METADATA, + ), + ) + + +def test_boto3_passthrough_calls_carry_the_key_identity_in_signed_request_metadata( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"who-is-calling-{uuid.uuid4().hex}" + forged: Final = {"user_api_key_team_alias": "forged", "ticket": "T-42", "bad!": "x"} + runtime: Final = rig.runtime(rig.key, monkeypatch) + model: Final = rig.dropping_model + + converse: Final = runtime.converse(modelId=model, messages=_messages(marker), requestMetadata=forged) + converse_stream: Final = [ + next(iter(event)) + for event in runtime.converse_stream(modelId=model, messages=_messages(marker), requestMetadata=forged)[ + "stream" + ] + ] + invoke: Final = json.loads( + runtime.invoke_model(modelId=model, body=json.dumps(_invoke_body(marker)))["body"].read() + ) + invoke_stream: Final = [ + json.loads(event["chunk"]["bytes"]) + for event in runtime.invoke_model_with_response_stream(modelId=model, body=json.dumps(_invoke_body(marker)))[ + "body" + ] + ] + assert converse["output"] == _CONVERSE_REPLY["output"], converse + assert converse_stream == ["messageStart", "messageStop", "metadata"], converse_stream + assert invoke == _INVOKE_REPLY, invoke + assert invoke_stream == [_INVOKE_REPLY], invoke_stream + + received: Final = rig.received(marker) + assert [request.target for request in received] == [ + f"/model/{_MODEL_ID}/converse", + f"/model/{_MODEL_ID}/converse-stream", + f"/model/{_MODEL_ID}/invoke", + f"/model/{_MODEL_ID}/invoke-with-response-stream", + ], received + converse_requests: Final = received[:2] + invoke_requests: Final = received[2:] + assert [json.loads(request.body) for request in converse_requests] == [ + {"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}}, + {"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}}, + ], [request.body for request in converse_requests] + assert [_METADATA_HEADER in request.headers for request in converse_requests] == [False, False] + assert [json.loads(request.body) for request in invoke_requests] == [_invoke_body(marker), _invoke_body(marker)] + assert [json.loads(request.headers[_METADATA_HEADER]) for request in invoke_requests] == [ + rig.identity, + rig.identity, + ], [request.headers for request in invoke_requests] + assert all(_METADATA_HEADER in _signed_header_names(request) for request in invoke_requests), [ + request.headers["authorization"] for request in invoke_requests + ] + + +def test_invoke_metadata_too_long_for_the_header_signs_the_leading_pairs_that_fit( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"long-key-{uuid.uuid4().hex}" + runtime: Final = rig.runtime(rig.long_key, monkeypatch) + + invoke: Final = json.loads( + runtime.invoke_model(modelId=rig.strict_model, body=json.dumps(_invoke_body(marker)))["body"].read() + ) + + assert invoke == _INVOKE_REPLY, invoke + (request,) = rig.received(marker) + assert json.loads(request.body) == _invoke_body(marker) + header: Final = request.headers[_METADATA_HEADER] + sent_pairs: Final = len(json.loads(header)) + assert tuple(json.loads(header).items()) == rig.long_identity[:sent_pairs], header + assert 2 <= sent_pairs < len(rig.long_identity), header + assert len(header) <= _HEADER_MAX_CHARS < len(_header_value(rig.long_identity[: sent_pairs + 1])) + assert _METADATA_HEADER in _signed_header_names(request), request.headers["authorization"] + + +@pytest.mark.parametrize( + ("action", "caller_metadata", "error_code"), + [ + pytest.param("converse", {"ticket": 42}, "SerializationException", id="converse-number-value"), + pytest.param("converse-stream", {"bad!": "x"}, "ValidationException", id="converse-stream-bad-key"), + pytest.param("invoke", {"bad!": "x"}, "ValidationException", id="invoke-bad-key"), + ], +) +def test_malformed_caller_request_metadata_reaches_boto3_as_the_bedrock_error_without_an_upstream_call( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str, caller_metadata: Mapping[str, JsonValue], error_code: str +) -> None: + marker: Final = f"malformed-{uuid.uuid4().hex}" + runtime: Final = rig.runtime(rig.key, monkeypatch) + + with pytest.raises(ClientError) as raised: + _send_with_caller_metadata(runtime, action, rig.strict_model, marker, caller_metadata) + + assert raised.value.response["Error"]["Code"] == error_code, raised.value.response + assert raised.value.response["ResponseMetadata"]["HTTPStatusCode"] == 400, raised.value.response + assert rig.received(marker) == () + + +def test_chat_on_a_dropping_converse_deployment_sends_only_the_valid_caller_pairs(rig: Rig) -> None: + marker: Final = f"chat-{uuid.uuid4().hex}" + + response: Final = rig.proxy.client.post( + "/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": marker}], + "requestMetadata": {"ticket": "T-42", "bad!": "x"}, + }, + headers={"Authorization": f"Bearer {rig.key}"}, + timeout=case_timeout_seconds(), + ) + + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "ok", response.text + (request,) = rig.received(marker) + assert unquote(request.target) == f"/model/{_MODEL_ID}/converse", request.target + assert json.loads(request.body)["requestMetadata"] == {**rig.identity, "ticket": "T-42"}, request.body diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 87cd3105cc7..ea29914ba65 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -11,9 +11,9 @@ import pytest import yaml from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows +from integration._support.ports import allocate_port from integration._support.process import owned_proxy_process from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario -from integration._support.ports import allocate_port from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import JsonValue, TypeAdapter @@ -21,6 +21,13 @@ import litellm _API_KEY: Final = "synthetic-decisions-key" _ENV_KEY: Final = "synthetic-decisions-env-key" +_AZURE_AI_ENTRA_SCENARIO: Final = "providers" +_AZURE_AI_ENTRA_TOKEN: Final = "synthetic-decisions-entra-token" +_AZURE_AI_MISSING_CREDENTIALS_MESSAGE: Final = ( + "Missing Azure AI credentials - set an API key (`api_key` or AZURE_AI_API_KEY), or Entra ID / OAuth " + "credentials (`tenant_id` + `client_id` + `client_secret`, `azure_ad_token`, an OIDC token, or a managed " + "identity with `litellm.enable_azure_ad_token_refresh = True`)" +) _PASS_THROUGH_MODEL: Final = "gpt-6-luna" _PASS_THROUGH_AUTHORIZATION: Final = "Bearer customer-held-upstream-key" _PASS_THROUGH_NEIGHBOUR: Final = "decisions-beside-a-pass-through" @@ -83,6 +90,7 @@ _PROVIDERS: Final = ( "perplexity/pplx-decider-v1-27b", ), _Provider("typesafe", "typesafe/jev-1.13.0", "/v1/systemone", "jev-1.13.0", _API_KEY, False, "typesafe/jev-1.13.0"), + _Provider("nace", "nace/drex-v2.0", "/v1/systemone", "drex-v2.0", _API_KEY, False, "nace/drex-v2.0"), _Provider( "openrouter", "openrouter/typesafe/jev-1.13", @@ -100,9 +108,7 @@ _PROVIDERS: Final = ( _Provider( "bespoke", "bespoke/nimble-latest", "/v1/systemone", "nimble-latest", None, False, "bespoke/nimble-latest" ), - _Provider( - "hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None - ), + _Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None), _Provider( "fireworks_ai", "fireworks_ai/accounts/fireworks/routers/rune", @@ -135,7 +141,7 @@ _PROVIDERS: Final = ( ), ) _PERPLEXITY: Final = _PROVIDERS[0] -_OPENROUTER: Final = _PROVIDERS[2] +_OPENROUTER: Final = next(provider for provider in _PROVIDERS if provider.name == "openrouter") _OPENROUTER_CHAT_MODEL: Final = "openrouter/openai/gpt-5-mini" _UNSUPPORTED_PROVIDER_MODEL: Final = "anthropic/claude-opus-5-5" _CONNECTION_ERROR: Final = "litellm.APIConnectionError" @@ -235,6 +241,56 @@ def _assert_connection_error(response: httpx.Response) -> None: assert _GENERIC_API_ERROR not in response.text, response.text +def test_azure_ai_entra_auth_reaches_the_scripted_upstream_for_both_decisions_routes( + gateway: Gateway, scenario: Scenario +) -> None: + provider: Final = next(provider for provider in _PROVIDERS if provider.name == "azure_ai") + handle: Final = register_scenario( + _AZURE_AI_ENTRA_SCENARIO, + JsonResponse(content_type="application/json", body=_answer_body(provider)), + ) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = "azure_ai/decision-1-entra" + body: Final = {"model": model, "state": _STATE, "questions": _QUESTIONS} + decisions_body: Final = { + "model": model, + "input": "The export job hangs at 99%", + "questions": [{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}], + } + + no_credentials_body: Final = {**body, "model": "azure_ai/decision-1-no-credentials"} + systemone_auth_error: Final = gateway.request("POST", "/v1/systemone", no_credentials_body) + assert systemone_auth_error.status_code == 401, systemone_auth_error.text + assert _AZURE_AI_MISSING_CREDENTIALS_MESSAGE in string_value( + object_value(object_value(systemone_auth_error.json())["error"])["message"] + ) + no_credentials_decisions_body: Final = { + **decisions_body, + "model": "azure_ai/decision-1-no-credentials", + } + decisions_auth_error: Final = gateway.request("POST", "/v1/decisions", no_credentials_decisions_body) + assert decisions_auth_error.status_code == 401, decisions_auth_error.text + assert _AZURE_AI_MISSING_CREDENTIALS_MESSAGE in string_value( + object_value(object_value(decisions_auth_error.json())["error"])["message"] + ) + assert not _upstream_calls(gateway, handle) + + systemone_response: Final = gateway.request("POST", "/v1/systemone", body) + assert systemone_response.status_code == 200, systemone_response.text + decisions_response: Final = gateway.request("POST", "/v1/decisions", decisions_body) + assert decisions_response.status_code == 200, decisions_response.text + + calls: Final = _upstream_calls(gateway, handle) + assert tuple(string_value(call["authorization"]) for call in calls) == ( + f"Bearer {_AZURE_AI_ENTRA_TOKEN}", + f"Bearer {_AZURE_AI_ENTRA_TOKEN}", + ) + assert tuple(string_value(call["path"]) for call in calls) == ( + "/providers/microsoft/v1/systemone", + "/providers/microsoft/v1/systemone", + ) + + def _pass_through_config(directory: Path, pass_through_target: str, native_api_base: str) -> Path: base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = { @@ -581,14 +637,16 @@ async def test_sdk_openrouter_connection_failures_raise_a_connection_error(gatew api_base=dead_api_base, api_key=_API_KEY, ) - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.APIConnectionError) as decisions_error: litellm.decisions( model=_OPENROUTER.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=dead_api_base, api_key=_API_KEY ) - with pytest.raises(litellm.APIConnectionError): + assert decisions_error.value.llm_provider == "openrouter" + with pytest.raises(litellm.APIConnectionError) as systemone_error: await litellm.asystemone( model=_OPENROUTER.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=dead_api_base, api_key=_API_KEY ) + assert systemone_error.value.llm_provider == "openrouter" def test_a_deployment_whose_provider_has_no_decisions_support_is_refused_naming_every_supported_provider( diff --git a/tests/integration/providers/test_google_genai_vertex_routes_wire.py b/tests/integration/providers/test_google_genai_vertex_routes_wire.py index d2c2f3bafb2..d7c2aa72c74 100644 --- a/tests/integration/providers/test_google_genai_vertex_routes_wire.py +++ b/tests/integration/providers/test_google_genai_vertex_routes_wire.py @@ -388,7 +388,7 @@ def test_adapter_route_upstream_error_reaches_the_caller_with_its_status( gateway: Gateway, status: int, stream: bool ) -> None: route: Final = _ROUTES["model_garden_publisher"] - message: Final = f"scripted vertex rejection {uuid.uuid4().hex}" + message: Final = f"scripted vertex rejection {uuid.uuid4().hex}f4039" with wire_server(_vertex_peer(status=status, error=message)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire, route, num_retries=0) key: Final = scenario.key(models=[model]) diff --git a/tests/integration/providers/test_openai_like_response_headers_wire.py b/tests/integration/providers/test_openai_like_response_headers_wire.py new file mode 100644 index 00000000000..f4d3965fb0f --- /dev/null +++ b/tests/integration/providers/test_openai_like_response_headers_wire.py @@ -0,0 +1,450 @@ +import itertools +import json +import math +import uuid +from collections.abc import Callable, Mapping, Sequence +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.prometheus_series import Sample, scrape, series_cap_config +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Deploy: TypeAlias = Callable[[Gateway, Scenario, Wire], str] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PROJECT: Final = "scripted-project" +_REPLY_TEXT: Final = "header probe reply" +_USAGE: Final[JsonValue] = {"prompt_tokens": 13, "completion_tokens": 4, "total_tokens": 17} +_UPSTREAM_COST: Final = "0.0421" +_SPEND_SELECT: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s' +_KEY_SPEND_SELECT: Final = 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s' +_PATHS: Final[Mapping[Endpoint, str]] = MappingProxyType( + {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"} +) +_RATE_LIMIT_HEADERS: Final[Mapping[str, str]] = MappingProxyType( + {"x-ratelimit-remaining-requests": "12.5", "x-ratelimit-remaining-tokens": "856, abc"} +) +_METERED_WORKERS: Final = 2 +_METERED_CALLS: Final[tuple[tuple[Endpoint, str, bool], ...]] = ( + ("chat", "groq", False), + ("chat", "groq", True), + ("chat", "openai", False), + ("chat", "openai", True), + ("responses", "groq", False), + ("responses", "groq", True), + ("messages", "groq", False), + ("messages", "groq", True), +) +_LATIN1_CALLS: Final[tuple[tuple[Endpoint, str, bool], ...]] = ( + ("chat", "bedrock_mantle", False), + ("chat", "bedrock_mantle", True), + ("responses", "groq", False), + ("responses", "groq", True), + ("messages", "groq", False), + ("messages", "groq", True), +) + + +def _keyed(model: str, **extra: JsonValue) -> Deploy: + def create(_gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model(model=model, api_base=wire.url, api_key="synthetic-provider-key", **extra) + + return create + + +def _vertex(model: str, location: str) -> Deploy: + def create(gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"vertex_ai/{model}", + api_base=wire.url, + api_key=None, + vertex_project=_PROJECT, + vertex_location=location, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + return create + + +_MEMBERS: Final[Mapping[str, Deploy]] = MappingProxyType( + { + "bedrock_mantle": _keyed("bedrock_mantle/google.gemma-4-31b", aws_region_name="us-west-2"), + "groq": _keyed("groq/openai/gpt-oss-20b"), + "openai": _keyed("openai/gpt-5.4-mini"), + "watsonx": _keyed("watsonx/ibm/granite-3-8b-instruct", zen_api_key="synthetic-zen-key", project_id=_PROJECT), + "vertex_model_garden": _vertex("xai/grok-4.20-non-reasoning", "global"), + "maritalk": _keyed("maritalk/sabia-3"), + } +) + + +def _reply_json() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": "scripted", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _REPLY_TEXT}, "finish_reason": "stop"} + ], + "usage": _USAGE, + "service_tier": "on_demand", + } + ).encode() + + +def _chunk(response_id: str, choices: list[JsonValue], **extra: JsonValue) -> bytes: + body: Final = { + "id": response_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "scripted", + "choices": choices, + **extra, + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def _frames() -> tuple[bytes, ...]: + response_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return ( + _chunk( + response_id, [{"index": 0, "delta": {"role": "assistant", "content": _REPLY_TEXT}, "finish_reason": None}] + ), + _chunk(response_id, [{"index": 0, "delta": {}, "finish_reason": "stop"}]), + _chunk(response_id, [], usage=_USAGE), + b"data: [DONE]\n\n", + ) + + +def _peer(headers: Mapping[str, str]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_frames(), headers=headers) + return Reply(body=_reply_json(), headers=headers) + + return respond + + +def _request_id_headers(request_id: str) -> Mapping[str, str]: + return MappingProxyType({"x-request-id": request_id, "x-amzn-requestid": request_id}) + + +def _body(endpoint: Endpoint, model: str, stream: bool) -> dict[str, JsonValue]: + prompt: Final = f"header probe {uuid.uuid4().hex}" + match endpoint: + case "chat": + return { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + "cache": {"no-cache": True}, + **({"stream_options": {"include_usage": True}} if stream else {}), + } + case "responses": + return {"model": model, "input": prompt, "stream": stream} + case "messages": + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + } + + +def _send(gateway: Gateway, endpoint: Endpoint, model: str, key: str, *, stream: bool = False) -> httpx.Response: + return gateway.request("POST", _PATHS[endpoint], _body(endpoint, model, stream), key=key) + + +class _Reply(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _Text(_Reply): + text: str = "" + + +class _ChatMessage(_Reply): + content: str = "" + + +class _ChatChoice(_Reply): + message: _ChatMessage + + +class _ChatReply(_Reply): + choices: tuple[_ChatChoice, ...] + + +class _ResponsesItem(_Reply): + content: tuple[_Text, ...] = () + + +class _ResponsesReply(_Reply): + output: tuple[_ResponsesItem, ...] + + +class _MessagesReply(_Reply): + content: tuple[_Text, ...] + + +def _text(endpoint: Endpoint, response: httpx.Response) -> str: + match endpoint: + case "chat": + return _ChatReply.model_validate_json(response.content).choices[0].message.content + case "responses": + items: Final = _ResponsesReply.model_validate_json(response.content).output + return "".join(part.text for part in itertools.chain.from_iterable(item.content for item in items)) + case "messages": + return "".join(block.text for block in _MessagesReply.model_validate_json(response.content).content) + + +def _header_cost(response: httpx.Response) -> float: + values: Final = response.headers.get_list("x-litellm-response-cost") + return float(values[0]) if values else math.nan + + +def _llm_provider_request_ids(response: httpx.Response) -> dict[str, str | None]: + return {name: response.headers.get(name) for name in ("llm_provider-x-request-id", "llm_provider-x-amzn-requestid")} + + +def _spend(key: str) -> float: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: read_rows(_SPEND_SELECT, (digest,)), lambda rows: len(rows) >= 1, seconds=70) + rows: Final = held(lambda: read_rows(_SPEND_SELECT, (digest,)), lambda found: len(found) == 1, holding=3) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float)), rows + return float(spend) + + +def _model_values(samples: Sequence[Sample], name: str, model: str) -> tuple[float, ...]: + return tuple( + sample.value for sample in samples if sample.name == name and sample.labels.get("model_group") == model + ) + + +def _deployment_metrics( + samples: Sequence[Sample], models: Mapping[str, str] +) -> dict[str, tuple[float, tuple[float, ...], tuple[float, ...]]]: + return { + call: ( + sum(_model_values(samples, "litellm_deployment_success_responses_total", model)), + _model_values(samples, "litellm_remaining_requests_metric", model), + _model_values(samples, "litellm_remaining_tokens_metric", model), + ) + for call, model in models.items() + } + + +def _key_spend(key: str) -> float | None: + rows: Final = read_rows(_KEY_SPEND_SELECT, (sha256(key.encode()).hexdigest(),)) + spend: Final = rows[0]["spend"] if rows else None + return float(spend) if isinstance(spend, (int, float)) else None + + +@pytest.mark.parametrize( + ("endpoint", "member"), + [ + ("chat", "bedrock_mantle"), + ("chat", "groq"), + ("chat", "watsonx"), + ("chat", "vertex_model_garden"), + ("chat", "maritalk"), + ("responses", "groq"), + ], +) +def test_non_stream_returns_the_provider_request_id(gateway: Gateway, endpoint: Endpoint, member: str) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + with wire_server(_peer(_request_id_headers(request_id))) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS[member](gateway, scenario, wire) + response: Final = _send(gateway, endpoint, model, scenario.key(models=[model])) + assert response.status_code == 200, response.text + assert _text(endpoint, response) == _REPLY_TEXT, response.text + assert _llm_provider_request_ids(response) == { + "llm_provider-x-request-id": request_id, + "llm_provider-x-amzn-requestid": request_id, + }, dict(response.headers) + assert len(wire.drain()) == 1 + + +def test_stream_chat_returns_the_provider_request_id(gateway: Gateway) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + with wire_server(_peer(_request_id_headers(request_id))) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS["bedrock_mantle"](gateway, scenario, wire) + response: Final = _send(gateway, "chat", model, scenario.key(models=[model]), stream=True) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert _llm_provider_request_ids(response) == { + "llm_provider-x-request-id": request_id, + "llm_provider-x-amzn-requestid": request_id, + }, dict(response.headers) + assert len(wire.drain()) == 1 + + +def test_upstream_cannot_overwrite_litellm_owned_headers(gateway: Gateway) -> None: + spoofed: Final = { + "x-litellm-call-id": "spoofed-call-id", + "x-litellm-model-id": "spoofed-model-id", + "x-litellm-attempted-fallbacks": "7", + } + with wire_server(_peer(spoofed)) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS["groq"](gateway, scenario, wire) + response: Final = _send(gateway, "chat", model, scenario.key(models=[model])) + assert response.status_code == 200, response.text + assert response.headers["x-litellm-call-id"] != spoofed["x-litellm-call-id"], dict(response.headers) + assert response.headers["x-litellm-model-id"] != spoofed["x-litellm-model-id"], dict(response.headers) + assert response.headers["x-litellm-attempted-fallbacks"] == "0", dict(response.headers) + assert len(wire.drain()) == 1 + + +@pytest.mark.parametrize(("endpoint", "member"), [("chat", "bedrock_mantle"), ("responses", "groq")]) +def test_non_stream_spend_matches_stream_spend_under_an_upstream_cost_header( + gateway: Gateway, endpoint: Endpoint, member: str +) -> None: + with wire_server(_peer({"x-litellm-response-cost": _UPSTREAM_COST})) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS[member](gateway, scenario, wire) + stream_key: Final = scenario.key(models=[model]) + non_stream_key: Final = scenario.key(models=[model]) + streamed: Final = _send(gateway, endpoint, model, stream_key, stream=True) + assert streamed.status_code == 200, streamed.text + assert _REPLY_TEXT in streamed.text, streamed.text + whole: Final = _send(gateway, endpoint, model, non_stream_key) + assert whole.status_code == 200, whole.text + assert _text(endpoint, whole) == _REPLY_TEXT, whole.text + stream_spend: Final = _spend(stream_key) + assert stream_spend > 0 + non_stream_spend: Final = _spend(non_stream_key) + assert math.isclose(non_stream_spend, stream_spend), (non_stream_spend, stream_spend) + assert len(wire.drain()) == 2 + + +@pytest.mark.parametrize("member", ["openai", "groq"]) +@pytest.mark.parametrize("upstream_cost", ["-5", "abc", "nan", "inf"], ids=["negative", "garbage", "nan", "infinite"]) +def test_an_invalid_upstream_cost_header_bills_the_price_map(gateway: Gateway, member: str, upstream_cost: str) -> None: + with ( + wire_server(_peer({"x-litellm-response-cost": upstream_cost})) as invalid_wire, + wire_server(_peer({})) as plain_wire, + gateway.scenario() as scenario, + ): + invalid_model: Final = _MEMBERS[member](gateway, scenario, invalid_wire) + plain_model: Final = _MEMBERS[member](gateway, scenario, plain_wire) + budget_key: Final = scenario.key(models=[invalid_model], max_budget=1) + reference_key: Final = scenario.key(models=[plain_model]) + billed: Final = _send(gateway, "chat", invalid_model, budget_key) + reference: Final = _send(gateway, "chat", plain_model, reference_key) + assert (billed.status_code, reference.status_code) == (200, 200), (billed.text, reference.text) + reference_spend: Final = _spend(reference_key) + assert reference_spend > 0 + key_spend: Final = eventually( + lambda: _key_spend(budget_key), + lambda spend: spend is not None and math.isclose(spend, reference_spend), + 70, + return_last_on_timeout=True, + ) + billed_costs: Final = { + "response_header": _header_cost(billed), + "spend_log": _spend(budget_key), + "key": key_spend if key_spend is not None else math.nan, + } + assert {name: math.isclose(cost, reference_spend) for name, cost in billed_costs.items()} == dict.fromkeys( + billed_costs, True + ), (billed_costs, reference_spend) + followup: Final = _send(gateway, "chat", invalid_model, budget_key) + assert followup.status_code == 200, followup.text + assert (len(invalid_wire.drain()), len(plain_wire.drain())) == (2, 1) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_provider_header_outside_latin1_does_not_fail_the_call(gateway: Gateway, tmp_path: Path) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + utf8_bytes_as_wire_text: Final = "東京リージョン".encode().decode("latin-1") + headers: Final = {**_request_id_headers(request_id), "x-provider-region": utf8_bytes_as_wire_text} + with ( + wire_server(_peer(headers)) as wire, + owned_proxy(gateway, tmp_path, {"DISABLE_AIOHTTP_TRANSPORT": "True"}) as httpx_transport, + httpx_transport.scenario() as scenario, + ): + models: Final = { + member: _MEMBERS[member](httpx_transport, scenario, wire) for member in ("bedrock_mantle", "groq") + } + key: Final = scenario.key(models=list(models.values())) + replies: Final = { + f"{endpoint}-{'stream' if stream else 'whole'}": _send( + httpx_transport, endpoint, models[member], key, stream=stream + ) + for endpoint, member, stream in _LATIN1_CALLS + } + assert {call: reply.status_code for call, reply in replies.items()} == dict.fromkeys(replies, 200), { + call: reply.text for call, reply in replies.items() + } + assert {call: reply.headers.get("llm_provider-x-provider-region") for call, reply in replies.items()} == ( + dict.fromkeys(replies, None) + ) + assert {call: _REPLY_TEXT in reply.text for call, reply in replies.items()} == dict.fromkeys(replies, True) + assert {endpoint: _text(endpoint, replies[f"{endpoint}-whole"]) for endpoint in _PATHS} == dict.fromkeys( + _PATHS, _REPLY_TEXT + ) + forwarding: Final = ("chat-whole", "chat-stream", "responses-whole") + assert {call: _llm_provider_request_ids(replies[call]) for call in forwarding} == dict.fromkeys( + forwarding, {"llm_provider-x-request-id": request_id, "llm_provider-x-amzn-requestid": request_id} + ) + assert len(wire.drain()) == len(_LATIN1_CALLS) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + 120) +def test_rate_limit_headers_set_the_remaining_gauges_and_keep_the_deployment_metrics( + gateway: Gateway, tmp_path: Path +) -> None: + multiproc_dir: Final = tmp_path / "prometheus" + multiproc_dir.mkdir() + with ( + wire_server(_peer(_RATE_LIMIT_HEADERS)) as wire, + owned_proxy( + gateway, + tmp_path, + {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)}, + config=series_cap_config(tmp_path, {}), + workers=_METERED_WORKERS, + ) as metered, + metered.scenario() as scenario, + ): + calls: Final = { + f"{endpoint}-{member}-{'stream' if stream else 'whole'}": ( + endpoint, + _MEMBERS[member](metered, scenario, wire), + stream, + ) + for endpoint, member, stream in _METERED_CALLS + } + models: Final = {call: model for call, (_, model, _) in calls.items()} + for model in models.values(): + wait_for_worker_readiness(metered, model, 1, _METERED_WORKERS) + key: Final = scenario.key(models=list(models.values())) + statuses: Final = { + call: _send(metered, endpoint, model, key, stream=stream).status_code + for call, (endpoint, model, stream) in calls.items() + } + assert statuses == dict.fromkeys(calls, 200) + observed: Final = eventually( + lambda: _deployment_metrics(scrape(metered), models), + lambda metrics: all(successes >= 1 for successes, _, _ in metrics.values()), + 30, + return_last_on_timeout=True, + ) + assert observed == dict.fromkeys(calls, (1.0, (12.5,), ())), observed + assert len(wire.drain()) == len(calls) diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py new file mode 100644 index 00000000000..f90d68c1a63 --- /dev/null +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpcore +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + DANGEROUS_TOOL_USE, + SAFEGUARDS, + deliveries, + holding_peer, + messages_body, + new_marker, + owned_config, + rust_routed, + safeguarded, + started_workers, +) +from pydantic import TypeAdapter + +_CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_RESTART_TIMEOUT: Final = int(4 * graceful_stop_seconds() + 240) +_BURST: Final = 12 +_SECOND_WAVE: Final = 6 +_CLIENT_ADDRESS: Final = TypeAdapter(tuple[str, int]) + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=120, trust_env=False)) for _ in range(count) + ) + + +def _post(client: httpx.Client, key: str, marker: str) -> int: + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), list(SAFEGUARDS)) + response: Final = client.post("/v1/messages", json=body, headers={"Authorization": f"Bearer {key}"}) + assert response.status_code != 200 or rust_routed(response.headers), dict(response.headers) + return response.status_code + + +def _post_or_dropped(client: httpx.Client, key: str, marker: str) -> int | None: + try: + return _post(client, key, marker) + except httpx.TransportError: + return None + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + stream: Final[object] = response.extensions["network_stream"] # pyright: ignore[reportAny] # httpx types extensions as Any + assert isinstance(stream, httpcore.NetworkStream), stream + port: Final = _CLIENT_ADDRESS.validate_python(stream.get_extra_info("client_addr"))[1] + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_worker_sigkill_mid_safeguards_burst_leaves_the_replacement_adding_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(holding_peer(held, release))) + config: Final = owned_config(tmp_path, gateway, wire.url) + owned: Final = stack.enter_context( + owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) + ) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually(lambda: started_workers(owned.log), lambda pids: len(pids) == 2, seconds=30) + clients: Final = _clients(stack, str(proxy_url), _BURST) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + ports: Final = tuple(_local_port(client) for client in clients) + markers: Final = tuple(new_marker() for _ in clients) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, marker) + for client, marker in zip(clients, markers, strict=True) + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), ( + port, + result, + shares, + ) + second_wave: Final = tuple(new_marker() for _ in range(_SECOND_WAVE)) + for marker in second_wave: + assert _post(owned.gateway.client, gateway.key, marker) == 200 + delivered: Final = deliveries(wire) + assert tuple(sent.beta for sent in delivered) == (DANGEROUS_TOOL_USE,) * (len(clients) + len(second_wave)), ( + delivered + ) + seen: Final = tuple(sum(sent.marker == marker for sent in delivered) for marker in markers + second_wave) + assert seen == (1,) * len(seen), seen + assert owned.process.poll() is None + + +@pytest.mark.timeout(_RESTART_TIMEOUT) +def test_proxy_shutdown_mid_safeguards_burst_then_a_fresh_boot_keeps_adding_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(holding_peer(held, release))) + config: Final = owned_config(tmp_path, gateway, wire.url) + markers: Final = tuple(new_marker() for _ in range(_BURST)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(markers))) + stack.callback(release.set) + with owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) as first: + clients: Final = _clients(stack, str(first.gateway.client.base_url), len(markers)) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, marker) + for client, marker in zip(clients, markers, strict=True) + ) + eventually(held.qsize, lambda size: size == len(markers), seconds=30) + first.begin_shutdown() + release.set() + results: Final = tuple(future.result(timeout=120) for future in futures) + during: Final = deliveries(wire) + assert all(result in (None, 200) for result in results), results + assert tuple(sent.beta for sent in during) == (DANGEROUS_TOOL_USE,) * len(markers), during + assert sorted(sent.marker for sent in during) == sorted(markers), during + with owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) as second: + after: Final = tuple(new_marker() for _ in range(_SECOND_WAVE)) + statuses: Final = tuple(_post(second.gateway.client, gateway.key, marker) for marker in after) + assert statuses == (200,) * len(after), statuses + delivered: Final = deliveries(wire) + assert tuple(sent.beta for sent in delivered) == (DANGEROUS_TOOL_USE,) * len(after), delivered + assert sorted(sent.marker for sent in delivered) == sorted(after), delivered diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py new file mode 100644 index 00000000000..75c5efc3aa0 --- /dev/null +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py @@ -0,0 +1,407 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, +) +from integration._support.database import read_rows +from integration._support.process import ( + OwnedProxy, + graceful_stop_seconds, + owned_proxy_process, +) +from integration._support.wire import Wire, wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + ANTHROPIC_KEY, + ANTHROPIC_TARGET, + BACKEND, + BASH_TOOL, + CLIENT_BETAS, + DANGEROUS_TOOL_USE, + DEEPSEEK_DEPLOYMENT, + DEEPSEEK_KEY, + DEEPSEEK_TARGET, + MERGED_BETAS, + NOT_A_LIST_MESSAGE, + OVERLOADED_MESSAGE, + REPLY_TEXT, + SAFEGUARD_RESULTS, + SAFEGUARDS, + VERTEX_DEPLOYMENT, + VERTEX_STREAM_TARGET, + VERTEX_TARGET, + WEB_SEARCH, + WEB_SEARCH_TOOL, + anthropic_peer, + chat_body, + deliveries, + json_object, + messages_body, + new_marker, + one_delivery, + outage_peer, + owned_config, + responses_body, + rust_routed, + safeguarded, + started_workers, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.xdist_group("anthropic_safeguards_beta") + +_CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_BURST: Final = 24 +_RECOVERY: Final = 6 +_LARGE_ENTRY: Final[JsonValue] = [ + { + "type": "dangerous_tool_use", + "classifier_context": {"v": 1, "permission_mode": "auto", "note": "x" * 5120}, + } +] + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + outage: threading.Event + + @property + def gateway(self) -> Gateway: + return self.owned.gateway + + +@dataclass(frozen=True, slots=True) +class _Answer: + path: str + marker: str + status: int + text: str + rust: bool + + +@pytest.fixture(scope="module") +def shared_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("anthropic-safeguards-beta") + outage: Final = threading.Event() + with ( + gateway_from_environment() as environment, + wire_server(outage_peer(outage)) as wire, + ): + config: Final = owned_config(directory, environment, wire.url) + with owned_proxy_process(environment, directory, {"LITELLM_RUST": "true"}, config=config, workers=2) as owned: + eventually( + lambda: started_workers(owned.log), + lambda pids: len(pids) == 2, + seconds=60, + ) + wire.drain() + yield _Rig(owned, wire, outage) + + +@pytest.fixture +def rig(shared_rig: _Rig) -> _Rig: + shared_rig.outage.clear() + shared_rig.wire.drain() + return shared_rig + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', + (call_id,), + ) + + +def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + body, + headers={} if beta is None else {"anthropic-beta": beta}, + ) + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_sdk_messages_with_safeguards_go_upstream_with_the_dangerous_tool_use_beta( + rig: _Rig, +) -> None: + marker: Final = new_marker() + with httpx.Client(trust_env=False, timeout=60) as http: + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=http, + ) + raw: Final = client.messages.with_raw_response.create( + model=ANTHROPIC_DEPLOYMENT, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, + ) + message: Final = raw.parse() + assert rust_routed(raw.headers), dict(raw.headers) + received: Final = json_object(message.model_dump_json()) + assert message.id == f"msg_{marker}", received + assert received.get("safeguard_results") == list(SAFEGUARD_RESULTS), received + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.api_key, sent.beta) == ( + ANTHROPIC_TARGET, + ANTHROPIC_KEY, + DANGEROUS_TOOL_USE, + ), sent + assert sent.safeguards == list(SAFEGUARDS), sent + assert "x-litellm-call-id" in raw.headers, dict(raw.headers) + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: _spend_rows(call_id), lambda found: len(found) == 1, seconds=70) + assert len(rows) == 1, rows + + +@pytest.mark.timeout(_CELL_TIMEOUT) +async def test_sdk_streamed_messages_with_safeguards_go_upstream_with_the_beta( + rig: _Rig, +) -> None: + marker: Final = new_marker() + async with httpx.AsyncClient(trust_env=False, timeout=60) as http: + client: Final = anthropic.AsyncAnthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=http, + ) + async with client.messages.stream( + model=ANTHROPIC_DEPLOYMENT, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert rust_routed(stream.response.headers), dict(stream.response.headers) + assert (final.id, text) == (f"msg_{marker}", REPLY_TEXT), (final, text) + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.beta, sent.safeguards) == ( + ANTHROPIC_TARGET, + DANGEROUS_TOOL_USE, + list(SAFEGUARDS), + ), sent + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_raw_messages_merge_the_safeguards_beta_into_the_client_betas(rig: _Rig, stream: bool) -> None: + marker: Final = new_marker() + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) + response: Final = _post(rig.gateway, body, beta=CLIENT_BETAS) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.beta) == (ANTHROPIC_TARGET, MERGED_BETAS), sent + + +@pytest.mark.parametrize("extra", [{"safeguards": None}, {}], ids=["null", "absent"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_null_or_absent_safeguards_add_no_beta(rig: _Rig, extra: Mapping[str, JsonValue]) -> None: + marker: Final = new_marker() + response: Final = _post(rig.gateway, {**messages_body(ANTHROPIC_DEPLOYMENT, marker), **extra}) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" not in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.beta, sent.safeguards) == (None, None), sent + + +@pytest.mark.parametrize( + ("safeguards", "status", "fragment"), + [ + (7, 400, NOT_A_LIST_MESSAGE), + ("dangerous_tool_use", 400, NOT_A_LIST_MESSAGE), + ({}, 400, NOT_A_LIST_MESSAGE), + ([], 200, "safeguard_results"), + (_LARGE_ENTRY, 200, "safeguard_results"), + ], + ids=["int", "string", "object", "empty_list", "5kb_entry"], +) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_malformed_safeguards_still_add_the_beta_and_relay_the_vendor_verdict( + rig: _Rig, safeguards: JsonValue, status: int, fragment: str +) -> None: + marker: Final = new_marker() + response: Final = _post( + rig.gateway, + safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), safeguards), + ) + assert (response.status_code, fragment in response.text) == (status, True), response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.beta, sent.safeguards) == (DANGEROUS_TOOL_USE, safeguards), sent + control: Final = new_marker() + after: Final = _post(rig.gateway, messages_body(ANTHROPIC_DEPLOYMENT, control)) + assert after.status_code == 200, after.text + assert rust_routed(after.headers), dict(after.headers) + assert one_delivery(rig.wire, control).beta is None + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("last", ["null", "list"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_duplicated_safeguards_keys_resolve_to_the_last_value(rig: _Rig, last: str) -> None: + marker: Final = new_marker() + opening: Final = json.dumps(messages_body(ANTHROPIC_DEPLOYMENT, marker))[:-1] + null_field: Final = '"safeguards": null' + list_field: Final = f'"safeguards": {json.dumps(list(SAFEGUARDS))}' + ordered: Final = (null_field, list_field) if last == "list" else (list_field, null_field) + raw: Final = f"{opening}, {ordered[0]}, {ordered[1]}}}".encode() + response: Final = rig.gateway.client.post( + "/v1/messages", + content=raw, + headers={ + "Authorization": f"Bearer {rig.gateway.key}", + "content-type": "application/json", + }, + ) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert ("safeguard_results" in response.text) == (last == "list"), response.text + sent: Final = one_delivery(rig.wire, marker) + assert sent.beta == (DANGEROUS_TOOL_USE if last == "list" else None), sent + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_safeguards_request_with_a_bad_key_never_reaches_the_peer(rig: _Rig) -> None: + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, new_marker()), list(SAFEGUARDS)) + response: Final = rig.gateway.request("POST", "/v1/messages", body, key="sk-not-a-key") + assert response.status_code == 401, response.text + assert deliveries(rig.wire) == () + + +@pytest.mark.parametrize("web_search", [False, True], ids=["bash_only", "with_web_search"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_vertex_messages_with_safeguards_keep_the_beta_from_the_shared_helper( + rig: _Rig, stream: bool, web_search: bool +) -> None: + marker: Final = new_marker() + tools: Final[list[JsonValue]] = [ + dict(BASH_TOOL), + *((dict(WEB_SEARCH_TOOL),) if web_search else ()), + ] + body: Final = safeguarded(messages_body(VERTEX_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) + response: Final = _post(rig.gateway, {**body, "tools": tools}) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert sent.target == (VERTEX_STREAM_TARGET if stream else VERTEX_TARGET), sent + assert sent.authorization == "Bearer scripted-token", sent + assert sent.beta == (f"{DANGEROUS_TOOL_USE},{WEB_SEARCH}" if web_search else DANGEROUS_TOOL_USE), sent + + +def _beta_set(header: str | None) -> frozenset[str]: + return frozenset(part.strip() for part in (header or "").split(",") if part.strip()) + + +@pytest.mark.parametrize("beta", [None, CLIENT_BETAS], ids=["no-client-betas", "client-betas"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_deepseek_messages_forward_safeguards_without_adding_a_beta(rig: _Rig, beta: str | None) -> None: + marker: Final = new_marker() + response: Final = _post( + rig.gateway, + safeguarded(messages_body(DEEPSEEK_DEPLOYMENT, marker), list(SAFEGUARDS)), + beta=beta, + ) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.api_key, sent.safeguards) == (DEEPSEEK_TARGET, DEEPSEEK_KEY, list(SAFEGUARDS)), sent + assert _beta_set(sent.beta) <= _beta_set(beta) - {DANGEROUS_TOOL_USE}, sent + + +def test_python_route_adds_the_beta_for_safeguards(gateway: Gateway, scenario: Scenario) -> None: + with wire_server(anthropic_peer) as wire: + deployment: Final = scenario.model(model=f"anthropic/{BACKEND}", api_base=wire.url, api_key=ANTHROPIC_KEY) + marker: Final = new_marker() + response: Final = _post(gateway, safeguarded(messages_body(deployment, marker), list(SAFEGUARDS))) + assert response.status_code == 200, response.text + assert "safeguard_results" in response.text, response.text + (sent,) = deliveries(wire) + assert (sent.marker, sent.safeguards) == (marker, list(SAFEGUARDS)), sent + assert DANGEROUS_TOOL_USE in _beta_set(sent.beta), sent + assert not rust_routed(response.headers), dict(response.headers) + + +def _burst_request(index: int, marker: str) -> tuple[str, dict[str, JsonValue]]: + if index % 4 == 0: + return "/v1/messages", safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), list(SAFEGUARDS)) + if index % 4 == 1: + return "/v1/messages", safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=True), list(SAFEGUARDS)) + if index % 4 == 2: + return "/v1/chat/completions", chat_body(ANTHROPIC_DEPLOYMENT, marker) + return "/v1/responses", responses_body(ANTHROPIC_DEPLOYMENT, marker) + + +def _answer(gateway: Gateway, path: str, marker: str, body: Mapping[str, JsonValue]) -> _Answer: + response: Final = gateway.request("POST", path, body) + return _Answer(path, marker, response.status_code, response.text, rust_routed(response.headers)) + + +def _wave(pool: ThreadPoolExecutor, gateway: Gateway, markers: tuple[str, ...], offset: int) -> tuple[_Answer, ...]: + planned: Final = tuple((marker, *_burst_request(offset + index, marker)) for index, marker in enumerate(markers)) + submitted: Final = tuple(pool.submit(_answer, gateway, path, marker, body) for marker, path, body in planned) + return tuple(future.result(timeout=120) for future in submitted) + + +def _assert_served(answer: _Answer) -> None: + assert answer.status == 200, answer + assert REPLY_TEXT in answer.text, answer + assert answer.rust == (answer.path == "/v1/messages"), answer + if answer.path == "/v1/messages": + assert f"msg_{answer.marker}" in answer.text and "safeguard_results" in answer.text, answer + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_peer_outage_mid_safeguards_burst_relays_each_verdict_once_and_recovers( + rig: _Rig, +) -> None: + healthy: Final = tuple(new_marker() for _ in range(_BURST // 2)) + outage: Final = tuple(new_marker() for _ in range(_BURST // 2)) + recovery: Final = tuple(new_marker() for _ in range(_RECOVERY)) + with ThreadPoolExecutor(max_workers=8) as pool: + before: Final = _wave(pool, rig.gateway, healthy, 0) + rig.outage.set() + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + during: Final = _wave(pool, rig.gateway, outage, len(healthy)) + rig.outage.clear() + after: Final = _wave(pool, rig.gateway, recovery, len(healthy) + len(outage)) + for answer in (*before, *after): + _assert_served(answer) + for answer in during: + assert answer.status >= 500 and OVERLOADED_MESSAGE in answer.text, answer + markers: Final = (*healthy, *outage, *recovery) + safeguarded_markers: Final = frozenset( + marker for index, marker in enumerate(markers) if _burst_request(index, marker)[0] == "/v1/messages" + ) + delivered: Final = deliveries(rig.wire) + seen: Final = tuple(sum(sent.marker == marker for sent in delivered) for marker in markers) + assert seen == (1,) * len(seen), seen + expected: Final = { + marker: (True, DANGEROUS_TOOL_USE) if marker in safeguarded_markers else (False, None) for marker in markers + } + for sent in delivered: + assert (sent.safeguards is not None, sent.beta) == expected[sent.marker], sent diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 1809b8ce5f8..b2aba7fb100 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -37,8 +37,11 @@ _STRUCTURED_OUTPUTS_BETA: Final = "structured-outputs-2025-11-13" _TOOL_MODEL: Final = "claude-sonnet-5" _EFFORT_GATED_MODEL: Final = "claude-haiku-4-5" _NATIVE_MODEL: Final = "claude-opus-5-5" +_FORCED_TOOL_NATIVE_MODEL: Final = "claude-haiku-5-5" _UNMAPPED_MODEL: Final = "claude-scripted-9" _BEDROCK_INVOKE_MODEL: Final = "anthropic.claude-sonnet-5" +_BEDROCK_JSON_TOOL_MODEL: Final = "us.anthropic.claude-opus-4-8" +_THINKING_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" _PROJECT: Final = "scripted-project" _LOCATION: Final = "us-east5" _VERTEX_MODELS: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/" @@ -85,6 +88,10 @@ _VALUES_SCHEMA: Final[dict[str, JsonValue]] = { "properties": {"values": {"type": "object"}}, "required": ["values"], } +_VALUES_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, +} _ADDITIONAL_OBJECT_SCHEMA: Final[dict[str, JsonValue]] = { **_OPTIONAL_SCHEMA, "additionalProperties": {"type": "object"}, @@ -1429,6 +1436,154 @@ def test_vertex_responses_native_folds_output_format(gateway: Gateway) -> None: assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT +def _native_deployment(gateway: Gateway, scenario: Scenario, api_base: str, provider: str, backend: str) -> str: + match provider: + case "vertex_ai": + return _vertex_deployment(gateway, scenario, api_base, backend) + case "anthropic": + return scenario.model(model=f"anthropic/{backend}", api_base=api_base, api_key=_ANTHROPIC_KEY) + case "azure_ai": + return scenario.model(model=f"azure_ai/{backend}", api_base=api_base, api_key=_AZURE_KEY) + case _: + raise AssertionError(provider) + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +def test_native_model_json_schema_with_an_open_object_takes_the_json_tool_and_keeps_the_object_open( + gateway: Gateway, provider: str +) -> None: + marker, marker_hex = _marker("values") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _native_deployment(gateway, scenario, wire.url, provider, _FORCED_TOOL_NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_VALUES_RESPONSE_FORMAT) + ) + assert _chat_content(response) == _tool_input("values", marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert not _structured(sent.body), sent.body + assert _json_tool_schema(sent)["properties"] == _VALUES_SCHEMA["properties"], sent.body + + +@pytest.mark.parametrize( + ("provider", "backend", "extra"), + [ + ("anthropic", _NATIVE_MODEL, {}), + ("vertex_ai", _NATIVE_MODEL, {}), + ("anthropic", _FORCED_TOOL_NATIVE_MODEL, {"thinking": {"type": "adaptive"}}), + ("anthropic", _FORCED_TOOL_NATIVE_MODEL, {"tools": [_USER_TOOL_SPEC]}), + ("anthropic", _FORCED_TOOL_NATIVE_MODEL, {"reasoning_effort": "none"}), + ("azure_ai", _FORCED_TOOL_NATIVE_MODEL, {"reasoning_effort": "none"}), + ], + ids=[ + "forced_tool_refused-anthropic", + "forced_tool_refused-vertex_ai", + "thinking", + "caller_tools", + "effort_none-anthropic", + "effort_none-azure_ai", + ], +) +def test_open_object_schema_keeps_native_output_when_the_json_tool_cannot_be_forced( + gateway: Gateway, provider: str, backend: str, extra: dict[str, JsonValue] +) -> None: + marker, marker_hex = _marker("values") + expected_tools: Final[frozenset[str]] = frozenset({_USER_TOOL}) if "tools" in extra else frozenset() + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _native_deployment(gateway, scenario, wire.url, provider, backend) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(model, marker, response_format=_VALUES_RESPONSE_FORMAT), **extra}, + ) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _structured(sent.body), sent.body + assert _tool_names(sent.body) == expected_tools, sent.body + assert "tool_choice" not in sent.body, sent.body + assert ("thinking" in sent.body) is ("thinking" in extra), sent.body + + +def _thinking_then_json_tool(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + _, marker_hex = _marker_of(body) + assert _JSON_TOOL in _tool_names(body), body + thinking: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": f"thinking {marker_hex}", + "signature": _THINKING_SIGNATURE, + } + answer: Final = _tool_use(f"toolu_{marker_hex}", _JSON_TOOL, _payload(marker_hex)) + return Reply(body=json.dumps(_message(f"msg_{marker_hex}", [thinking, answer], "tool_use")).encode()) + + +def _json_tool_deployment(gateway: Gateway, scenario: Scenario, api_base: str, provider: str) -> str: + match provider: + case "vertex_ai": + return _vertex_deployment(gateway, scenario, api_base, _TOOL_MODEL) + case "bedrock_invoke": + return _bedrock_invoke_deployment(scenario, api_base, _BEDROCK_JSON_TOOL_MODEL) + case _: + raise AssertionError(provider) + + +@pytest.mark.parametrize("provider", ["vertex_ai", "bedrock_invoke"]) +def test_json_mode_answer_keeps_the_thinking_that_preceded_the_json_tool_call(gateway: Gateway, provider: str) -> None: + marker, marker_hex = _marker("plain") + with wire_server(_thinking_then_json_tool) as wire, gateway.scenario() as scenario: + model: Final = _json_tool_deployment(gateway, scenario, wire.url, provider) + response: Final = gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "reasoning_effort": "high"} + ) + message: Final = _chat_message(response) + assert json.loads(string_value(message["content"])) == _payload(marker_hex), response.text + assert message.get("reasoning_content") == f"thinking {marker_hex}", response.text + (block,) = tuple(object_value(item) for item in list_value(message.get("thinking_blocks") or [])) + assert (block["thinking"], block["signature"]) == (f"thinking {marker_hex}", _THINKING_SIGNATURE), block + sent: Final = _one_delivery(wire, marker_hex) + assert "thinking" in sent.body, sent.body + assert sent.body.get("tool_choice") != _FORCED_JSON_TOOL, sent.body + + +def test_vertex_chat_native_schema_keeps_min_items_one_and_closes_a_nullable_object(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "backup": {"type": ["object", "null"], "properties": {"name": {"type": "string"}}, "required": ["name"]}, + }, + "required": ["summary", "tags", "backup"], + "additionalProperties": False, + } + expected: Final[dict[str, JsonValue]] = { + **schema, + "properties": { + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "backup": { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, + }, + }, + } + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "tagged", "schema": schema}, + } + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=response_format) + ) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _folded_format(sent) == {"type": "json_schema", "schema": expected}, sent.body + assert "tools" not in sent.body, sent.body + + @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) def test_vertex_messages_folds_output_format_into_output_config(gateway: Gateway, stream: bool) -> None: marker, marker_hex = _marker("native") diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index f047d6e0b3e..4e2cab56d27 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -473,6 +473,13 @@ model_list: vertex_project: synthetic-vertex-project vertex_location: global vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' + - model_name: vertex_ai/claude-haiku-5-5@default + litellm_params: + model: vertex_ai/claude-haiku-5-5@default + api_base: http://127.0.0.1:8191 + vertex_project: synthetic-vertex-project + vertex_location: global + vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' - model_name: vertex_ai/claude-opus-5-5@default litellm_params: model: vertex_ai/claude-opus-5-5@default @@ -625,6 +632,19 @@ model_list: api_key: synthetic-azure-ai-key model_info: base_model: azure_ai/Microsoft-Decision-1 + - model_name: azure_ai/decision-1-entra + litellm_params: + model: azure_ai/decision-1 + api_base: os.environ/INTEGRATION_UPSTREAM_URL + azure_ad_token: synthetic-decisions-entra-token + model_info: + base_model: azure_ai/Microsoft-Decision-1 + - model_name: azure_ai/decision-1-no-credentials + litellm_params: + model: azure_ai/decision-1 + api_base: os.environ/INTEGRATION_UPSTREAM_URL + model_info: + base_model: azure_ai/Microsoft-Decision-1 general_settings: master_key: os.environ/LITELLM_MASTER_KEY database_url: os.environ/DATABASE_URL diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py index d4328cd0e31..1ca92505191 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -84,7 +84,7 @@ _FAMILIES: Final = ( _Family("pin-bedrock-converse-1h", "pin-bedrock-converse-1h", one_hour=True, wire_one_hour=True), _Family("pin-bedrock-converse-5m", "pin-bedrock-converse-5m", one_hour=False, wire_one_hour=False), _Family("pin-bedrock-invoke-1h", "pin-bedrock-invoke-1h", one_hour=True, wire_one_hour=True), - _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=False, wire_one_hour=False), + _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=True, wire_one_hour=True), _Family("pin-bedrock-arn", "pin-bedrock-arn", one_hour=False, wire_one_hour=False), _Family("pin-azure-ai", "pin-azure-ai", one_hour=True, wire_one_hour=True), _Family("pin-vertex", "pin-vertex", one_hour=True, wire_one_hour=True), @@ -449,19 +449,39 @@ def test_provider_family_pin_lifetime_agrees_with_the_ttl_the_wire_carries(rig: spend_row_lands(response_id) -def test_messages_through_bedrock_invoke_pins_an_hour_and_sends_the_hour(rig: _Rig) -> None: +@pytest.mark.parametrize("group", ["pin-bedrock-invoke-1h", "pin-bedrock-arn-base"]) +def test_messages_through_bedrock_pins_an_hour_and_sends_the_hour(rig: _Rig, group: str) -> None: before: Final = _before(rig) marker: Final = new_marker() response: Final = rig.owned.gateway.request( - "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": "pin-bedrock-invoke-1h"} + "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": group} ) assert response.status_code == 200, response.text _, pin = one_new_pin(rig.redis, before) - assert pin.model_id == "pin-bedrock-invoke-1h", pin + assert pin.model_id == group, pin assert lives_one_hour(pin), pin assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [True] +def test_response_cache_hit_on_an_application_profile_rewrites_the_pin_for_its_base_model_hour(rig: _Rig) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + body: Final = chat_body("pin-bedrock-arn-base", _first_turn(new_marker(), marker, ONE_HOUR_MARK)) + first, served = post_chat(rig.owned.gateway, body) + key, pin = one_new_pin(rig.redis, before) + assert served == pin.model_id == "pin-bedrock-arn-base", (served, pin) + assert lives_one_hour(pin), pin + assert [wire_carries_one_hour(request) for request in sent_with(rig.wires(), marker)] == [True] + spend_row_lands(first) + rig.redis.delete(key) + second, _ = post_chat(rig.owned.gateway, body) + assert sent_with(rig.wires(), marker) == () + assert second == first + rewritten: Final = eventually(lambda: pin_at(rig.redis, key), lambda found: found is not None, seconds=30) + assert rewritten is not None and rewritten.model_id == "pin-bedrock-arn-base", rewritten + assert lives_one_hour(rewritten), rewritten + + @pytest.mark.parametrize("ttl", _MALFORMED_TTLS, ids=["int", "2h", "upper-1H", "empty", "null", "list", "5kb"]) def test_malformed_ttl_answers_and_pins_five_minutes_without_a_wire_ttl(rig: _Rig, ttl: JsonValue) -> None: before: Final = _before(rig) diff --git a/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py new file mode 100644 index 00000000000..1c419e5b11c --- /dev/null +++ b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py @@ -0,0 +1,139 @@ +import json +import re +import uuid +from collections.abc import Callable, Iterator, Mapping +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import ModelResponse, Router +from litellm.utils import _invalidate_model_cost_lowercase_map # pyright: ignore[reportPrivateUsage] # required cache invalidation + +_OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" +_JSON_TOOL: Final = "json_tool_call" +_MARKER: Final = re.compile(r"probe ([0-9a-f]{32})") +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, +} +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) +_OPTIONS: Final[dict[str, JsonValue]] = { + "max_tokens": 2048, + "thinking": {"type": "adaptive"}, + "response_format": {"type": "json_schema", "json_schema": {"name": "city", "schema": _SCHEMA}}, + "num_retries": 0, +} + + +def _native(body: Mapping[str, JsonValue], invoke: bool) -> bool: + output_config: Final = body.get("output_config" if invoke else "outputConfig") + return isinstance(output_config, dict) and ("format" if invoke else "textFormat") in output_config + + +def _invoke_reply(answer: dict[str, JsonValue], native: bool) -> Reply: + block: Final[dict[str, JsonValue]] = ( + {"type": "text", "text": json.dumps(answer)} + if native + else {"type": "tool_use", "id": "toolu_scripted", "name": _JSON_TOOL, "input": answer} + ) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [block], + "stop_reason": "end_turn" if native else "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 7}, + } + return Reply(body=json.dumps(message).encode()) + + +def _converse_reply(answer: dict[str, JsonValue], native: bool) -> Reply: + block: Final[dict[str, JsonValue]] = ( + {"text": json.dumps(answer)} + if native + else {"toolUse": {"toolUseId": "tooluse_scripted", "name": _JSON_TOOL, "input": answer}} + ) + converse: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [block]}}, + "stopReason": "end_turn" if native else "tool_use", + "usage": {"inputTokens": 11, "outputTokens": 7, "totalTokens": 18}, + "metrics": {"latencyMs": 1}, + } + return Reply(body=json.dumps(converse).encode()) + + +def _bedrock(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + found: Final = _MARKER.search(json.dumps(body)) + assert found is not None, body + invoke: Final = "/invoke" in request.target + reply: Final[Callable[[dict[str, JsonValue], bool], Reply]] = _invoke_reply if invoke else _converse_reply + return reply({"city": found.group(1)}, _native(body, invoke)) + + +@pytest.fixture +def opus_4_8_opted_in_for_the_process(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "model_cost", _COST_MAP.validate_python(litellm.model_cost)) # pyright: ignore[reportUnknownMemberType] # the cost map is untyped + litellm.register_model( + {_OPUS_4_8: {"supports_bedrock_native_structured_output": True}}, persist_across_reloads=False + ) + yield + monkeypatch.undo() + _invalidate_model_cost_lowercase_map() + + +def _model(route: str) -> str: + return f"bedrock/{route}/{_OPUS_4_8}" + + +def _connection(wire: Wire) -> dict[str, JsonValue]: + return {"api_base": wire.url, "api_key": "synthetic-bedrock-bearer", "aws_region_name": "us-east-1"} + + +def _sent_native(wire: Wire, route: str, call: Callable[[list[dict[str, str]]], object]) -> bool: + marker_hex: Final = uuid.uuid4().hex + response: Final = call([{"role": "user", "content": f"Name a city. probe {marker_hex}"}]) + assert isinstance(response, ModelResponse), response + assert json.loads(str(response.choices[0].message.content)) == {"city": marker_hex}, response # pyright: ignore[reportAttributeAccessIssue] # choices also types streaming chunks + (request,) = wire.drain() + return _native(_JSON.validate_json(request.body), route == "invoke") + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +@pytest.mark.usefixtures("opus_4_8_opted_in_for_the_process") +def test_a_register_model_native_output_opt_in_reaches_every_call_but_a_deployment_that_opts_out(route: str) -> None: + with wire_server(_bedrock) as wire: + router: Final = Router( + model_list=[ + { + "model_name": "opted-out", + "litellm_params": {"model": _model(route), **_connection(wire)}, + "model_info": {"supports_bedrock_native_structured_output": False}, + }, + {"model_name": "inherits", "litellm_params": {"model": _model(route), **_connection(wire)}}, + ], + num_retries=0, + ) + assert _sent_native( + wire, + route, + lambda messages: router.completion(model="inherits", messages=messages, **_OPTIONS), # pyright: ignore[reportUnknownMemberType] # Router.completion leaves **kwargs untyped + ) + assert not _sent_native( + wire, + route, + lambda messages: router.completion(model="opted-out", messages=messages, **_OPTIONS), # pyright: ignore[reportUnknownMemberType] # Router.completion leaves **kwargs untyped + ) + assert _sent_native( + wire, + route, + lambda messages: litellm.completion(_model(route), messages, **_connection(wire), **_OPTIONS), + ) diff --git a/tests/integration/sdk/test_stream_error_api_base_sdk.py b/tests/integration/sdk/test_stream_error_api_base_sdk.py new file mode 100644 index 00000000000..922e518c7bf --- /dev/null +++ b/tests/integration/sdk/test_stream_error_api_base_sdk.py @@ -0,0 +1,69 @@ +import json +from collections.abc import AsyncIterable, AsyncIterator +from typing import Final + +import openai +from integration._support.anthropic_sse import dropping_reply +from integration._support.wire import Reply, Request, wire_server + +import litellm +from litellm.exceptions import MidStreamFallbackError + +_PARTIAL_TEXT: Final = "partial answer" + + +def _chunk(content: str) -> bytes: + body: Final = { + "id": "chatcmpl-sdk-cut", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}, "finish_reason": None}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _cut_after_text(request: Request) -> Reply: + assert request.method == "POST" and "/chat/completions" in request.target, request.target + return dropping_reply((_chunk(_PARTIAL_TEXT), _chunk(" never sent")), abort_after=1) + + +def _text(item: object) -> str: + if not isinstance(item, litellm.ModelResponseStream): + return "" + return "".join(choice.delta.content or "" for choice in item.choices) + + +async def _items_or_error(stream: AsyncIterable[object]) -> AsyncIterator[object]: + try: + async for item in stream: + yield item + except openai.APIError as error: + yield error + + +async def test_sdk_mid_stream_failure_debug_info_names_the_api_base_it_called() -> None: + with wire_server(_cut_after_text) as wire: + response: Final = await litellm.acompletion( + model="azure/gpt-4o-mini", + api_base=wire.url, + api_key="synthetic-azure-key", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + num_retries=0, + ) + assert isinstance(response, AsyncIterable), response + items: Final = [item async for item in _items_or_error(response)] + errors: Final = [item for item in items if isinstance(item, openai.APIError)] + assert len(errors) == 1 and items[-1] is errors[0], items + assert "".join(_text(item) for item in items) == _PARTIAL_TEXT, items + error: Final = errors[0] + mapped: Final = ( + error.original_exception + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None + else error + ) + debug_info: Final = getattr(mapped, "litellm_debug_info", None) + assert isinstance(debug_info, str), mapped + assert f"\nAPI Base: `{wire.url}`" in debug_info, debug_info diff --git a/tests/integration/spend/test_daily_rollup_retry.py b/tests/integration/spend/test_daily_rollup_retry.py index 5b48d841f75..9ed712cdd06 100644 --- a/tests/integration/spend/test_daily_rollup_retry.py +++ b/tests/integration/spend/test_daily_rollup_retry.py @@ -142,6 +142,7 @@ def test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_ac assert metrics == { "spend": pytest.approx(0.06), "flat_cost": pytest.approx(0.0), + "ptu_hours": pytest.approx(0.0), "prompt_tokens": 20, "completion_tokens": 20, "cache_read_input_tokens": 0, diff --git a/tests/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index 3fba69b39b7..9e8f456df49 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -109,7 +109,9 @@ class FallbackProxy: spare_port: int -def _fallback_config(directory: Path, primary_port: int, spare_port: int) -> Path: +def _fallback_config( + directory: Path, primary_port: int, spare_port: int, extra_router_settings: Mapping[str, JsonValue] +) -> Path: base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) deployments: Final = [ { @@ -122,23 +124,37 @@ def _fallback_config(directory: Path, primary_port: int, spare_port: int) -> Pat } for name, port in ((_PRIMARY, primary_port), (_SPARE, spare_port)) ] - router_settings: Final = {"num_retries": 0, "disable_cooldowns": True, "fallbacks": [{_PRIMARY: [_SPARE]}]} + router_settings: Final = { + "num_retries": 0, + "disable_cooldowns": True, + "fallbacks": [{_PRIMARY: [_SPARE]}], + **extra_router_settings, + } path: Final = directory / "bridged-fallbacks.yaml" path.write_text(yaml.safe_dump({**base, "model_list": deployments, "router_settings": router_settings})) return path -@pytest.fixture(scope="module") -def fallback_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: +def _fallback_proxy( + tmp_path_factory: pytest.TempPathFactory, extra_router_settings: Mapping[str, JsonValue] +) -> Iterator[FallbackProxy]: directory: Final = tmp_path_factory.mktemp("bridged-fallbacks") primary_port, spare_port = allocate_port(), allocate_port() - with ( - gateway_from_environment() as shared, - owned_proxy(shared, directory, {}, config=_fallback_config(directory, primary_port, spare_port)) as owned, - ): + config: Final = _fallback_config(directory, primary_port, spare_port, extra_router_settings) + with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, config=config) as owned: yield FallbackProxy(owned, primary_port, spare_port) +@pytest.fixture(scope="module") +def fallback_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: + yield from _fallback_proxy(tmp_path_factory, {}) + + +@pytest.fixture(scope="module") +def bad_request_opt_out_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: + yield from _fallback_proxy(tmp_path_factory, {"fallback_on_unsupported_params": False}) + + def test_bridged_in_stream_rate_limit_falls_back_to_the_healthy_deployment(fallback_proxy: FallbackProxy) -> None: identity: Final = "resp_" + uuid.uuid4().hex with ( @@ -189,9 +205,7 @@ def test_bridged_top_level_rate_limit_falls_back_to_the_healthy_deployment(fallb _INVALID_PROMPT: Final = "Invalid prompt: your prompt was flagged." - - -@pytest.mark.parametrize( +_INVALID_PROMPT_REJECTIONS: Final = pytest.mark.parametrize( "rejection", [ pytest.param( @@ -201,7 +215,10 @@ _INVALID_PROMPT: Final = "Invalid prompt: your prompt was flagged." pytest.param(openai_error_event("invalid_prompt", _INVALID_PROMPT), id="top_level"), ], ) -def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallback( + + +@_INVALID_PROMPT_REJECTIONS +def test_bridged_invalid_prompt_before_any_token_falls_back_to_the_healthy_deployment( fallback_proxy: FallbackProxy, rejection: Mapping[str, JsonValue] ) -> None: identity: Final = "resp_" + uuid.uuid4().hex @@ -214,6 +231,29 @@ def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallbac ) as spare, ): response: Final = _chat(fallback_proxy.gateway, chat_body(_PRIMARY, identity)) + assert response.status_code == 200, response.text + assert chat_content(data_frames(response.text)) == "fallback answer", response.text + assert len(primary.drain()) == 1 and len(spare.drain()) == 1 + row: Final = spend_row(response.headers["x-litellm-call-id"]) + assert row["status"] == "success" and row["model_group"] == _SPARE, row + + +@_INVALID_PROMPT_REJECTIONS +def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallback_when_bad_requests_opt_out( + bad_request_opt_out_proxy: FallbackProxy, rejection: Mapping[str, JsonValue] +) -> None: + identity: Final = "resp_" + uuid.uuid4().hex + with ( + wire_server( + serve((frame(created(identity)), frame(rejection)), AZURE_TARGET), + port=bad_request_opt_out_proxy.primary_port, + ) as primary, + wire_server( + serve(healthy_stream(identity, "fallback answer"), AZURE_TARGET), + port=bad_request_opt_out_proxy.spare_port, + ) as spare, + ): + response: Final = _chat(bad_request_opt_out_proxy.gateway, chat_body(_PRIMARY, identity)) assert response.status_code == 400, response.text message: Final = string_value(error_body(response)["message"]) assert message.startswith("litellm.BadRequestError: ") and _INVALID_PROMPT in message, message diff --git a/tests/integration/translation/chat_completions/bases/anthropic.py b/tests/integration/translation/chat_completions/bases/anthropic.py index 16e1de53975..3a148c64859 100644 --- a/tests/integration/translation/chat_completions/bases/anthropic.py +++ b/tests/integration/translation/chat_completions/bases/anthropic.py @@ -42,7 +42,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011CfgDzBnuPYc7TdysZFRpo", + "id": "msg_011Cg7v1TRmiNAjM2RRxcxF8", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -58,6 +58,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -93,6 +94,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -123,7 +125,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfgDzH3AcaSNXtvurLgAG", + "id": "msg_011Cg7v28wK6VaszM4tL4D6n", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -139,6 +141,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -174,6 +177,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -223,6 +227,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -272,6 +277,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -321,6 +327,7 @@ CLAUDE_OPUS_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -370,6 +377,7 @@ CLAUDE_OPUS_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -400,7 +408,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-5", - "id": "msg_011CfgDzPkj2HxBq8qGJsGSi", + "id": "msg_011Cg7v2SdRUNPAvGScTF7au", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -417,6 +425,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -452,6 +461,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -482,7 +492,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfgDzV9vZCkuQJt15n3Z1", + "id": "msg_011Cg7v1hPQMJ1QhaC8HaSfQ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -499,6 +509,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -534,6 +545,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -549,8 +561,16 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( "content-type": "application/json", }, expected_provider_request={**ANTHROPIC_MESSAGES_REQUEST, "model": "claude-sonnet-5-5"}, - mock_provider_response={**ANTHROPIC_MESSAGES_RESPONSE, "model": "claude-sonnet-5-5"}, - expected_litellm_response={**CHAT_COMPLETIONS_RESPONSE, "model": "anthropic/claude-sonnet-5-5"}, + mock_provider_response={ + **ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-sonnet-5-5", + "usage": {**object_value(ANTHROPIC_MESSAGES_RESPONSE["usage"]), "speed": "standard"}, + }, + expected_litellm_response={ + **CHAT_COMPLETIONS_RESPONSE, + "model": "anthropic/claude-sonnet-5-5", + "usage": {**object_value(CHAT_COMPLETIONS_RESPONSE["usage"]), "speed": "standard"}, + }, ) @@ -618,6 +638,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 149}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -641,7 +662,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_015ndrvU4VqqoMBzK35jt4TX", "type": "function", - }, + } ], "reasoning_content": "There's a mention of a chat_completions client I don't actually have access " "to, so I'll ignore that and stick to my available tools. I'll proceed by " @@ -658,7 +679,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "examine next.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], "provider_specific_fields": { "citations": None, @@ -672,11 +693,11 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "deciding what to examine next.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, - }, + } ], "usage": { "completion_tokens": 288, @@ -694,6 +715,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -757,6 +779,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 66}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -780,7 +803,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_01Ew2u9E9j2mkT6RbwSw13uR", "type": "function", - }, + } ], "reasoning_content": "I'll focus on reading the validation file first, around line 18, since that's " "what was specifically requested before looking at the tests.\n" @@ -793,7 +816,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "tests.\n" "\n", "signature": CLAUDE_OPUS_5_5_THINKING_SIGNATURE_0, - }, + } ], "provider_specific_fields": { "citations": None, @@ -805,11 +828,11 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "the tests.\n" "\n", "signature": CLAUDE_OPUS_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, - }, + } ], "usage": { "completion_tokens": 215, @@ -827,6 +850,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -869,6 +893,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -892,11 +917,11 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_01XMSjzRvHCE753Wx4dU2Vxo", "type": "function", - }, + } ], "provider_specific_fields": {"citations": None, "thinking_blocks": None}, }, - }, + } ], "usage": { "completion_tokens": 143, @@ -914,6 +939,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) diff --git a/tests/integration/translation/chat_completions/bases/azure_ai.py b/tests/integration/translation/chat_completions/bases/azure_ai.py index e3e8a84340b..89c75f3c597 100644 --- a/tests/integration/translation/chat_completions/bases/azure_ai.py +++ b/tests/integration/translation/chat_completions/bases/azure_ai.py @@ -43,7 +43,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011Cfjax6mBrNVhFGQopmVCp", + "id": "msg_011Cg7uyAiPY51HN4rSYzQW5", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -59,6 +59,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -94,6 +95,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -144,6 +146,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -193,6 +196,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -224,10 +228,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfjaxGKedCbjtb7HuAUBC", + "id": "msg_011Cg7uzQ1rtYRwCd9aeTyoF", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello!"}], + "content": [{"type": "text", "text": "Hello! How can I help you?"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -237,9 +241,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - "output_tokens": 5, + "output_tokens": 11, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -253,17 +258,17 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "finish_reason": "stop", "index": 0, "message": { - "content": "Hello!", + "content": "Hello! How can I help you?", "role": "assistant", "provider_specific_fields": {"citations": None, "thinking_blocks": None}, }, } ], "usage": { - "completion_tokens": 5, + "completion_tokens": 11, "prompt_tokens": 18, - "total_tokens": 23, - "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 5}, + "total_tokens": 29, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 11}, "prompt_tokens_details": { "cached_tokens": 0, "text_tokens": 18, @@ -275,6 +280,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -389,7 +395,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfjaxUiYCUsqMC4iqfc2U", + "id": "msg_011Cg7uyf9RNvQWkpJobt6JU", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -406,6 +412,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -441,6 +448,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -674,6 +682,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 167}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -697,7 +706,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_01M68cqREYA4WqXBUK8rqcEq", "type": "function", - }, + } ], "reasoning_content": "I need to read shop/validation.py next, keeping in mind edge cases like zero, " "negative, or non-integer quantities since the validation logic likely handles those. " @@ -709,7 +718,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "negative, or non-integer quantities since the validation logic likely handles those. " "I'll look at the function starting around line 18 to understand its scope.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], "provider_specific_fields": { "citations": None, @@ -720,11 +729,11 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "negative, or non-integer quantities since the validation logic likely handles those. " "I'll look at the function starting around line 18 to understand its scope.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, - }, + } ], "usage": { "completion_tokens": 324, @@ -742,6 +751,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -763,10 +773,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "model": "claude-fable-5-1", "id": "msg_011Cg37NcfzdXibtC3dWhVWi", "content": [ - { - "type": "text", - "text": "I'll read the validation code and the test together since they are independent.", - }, + {"type": "text", "text": "I'll read the validation code and the test together since they are independent."}, { "type": "tool_use", "id": "toolu_01LxkeKf9WX9BVpiWrvfKUyc", @@ -791,6 +798,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, expected_litellm_response={ @@ -827,7 +835,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( ], "provider_specific_fields": {"citations": None, "thinking_blocks": None}, }, - }, + } ], "usage": { "completion_tokens": 216, @@ -845,6 +853,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) diff --git a/tests/integration/translation/chat_completions/bases/vertex_ai.py b/tests/integration/translation/chat_completions/bases/vertex_ai.py index d1022da78c7..447566b0ee3 100644 --- a/tests/integration/translation/chat_completions/bases/vertex_ai.py +++ b/tests/integration/translation/chat_completions/bases/vertex_ai.py @@ -577,6 +577,154 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_THINKING_SIGNATURE: Final = ( + "CAQS8gcKEggSEAIYAjgBQgh0aGlua2luZxIMKxM+8GdyGfxkJrU6Ggw3AMSc5G/SaulUVJkiMJgPNRytmwmfF4bj7WugkKlRxBr+SRwcL9iQ" + "+S5uvmIi/D8thAXX+oVg67diUo3rxiqNB0KrfSB/DXD50OomDcy90ai3a+RMcwIy9wEMzOdOIFqQ24tUMm1YorcQLQpresEnXA9ItiMAzP+P" + "8NMJ1br7O3JHYLIi/12UzVr36AKPrMJQk86VBn5fVxo0xcphyl071Y30m3wO0INGtU+d37IXSrYCJzQdnUSKxekBK28fxsA4tmxx1BE/1Hfl" + "UTsq36+CE2MatE0ymL2pQA+povoJ1iVlB/PBZvkOvpRcQubSwm8RqJrBfmse8baU2UJot0VNNC7OprQ3wYx1j3BGmXOQCGkvCn6jd2lPLjLS" + "F1CjbWQs7X18gB69y2sUosm3aI+MnQNHkUqe8JtIaii42uu3Xe/JheeRTRrcsuAdMZJ76ID9wRZnf3/QCHaMyrYwHXf5SsoUjX3Ofs5MuQT0" + "ISf4T7q+QGf51Bxb4jdfF879h061ZSejzknCHI3fV3SwuldjIIV+FD50d63iLqM/D8GOjaGfDrYugkkdBjCNEQkTpn9NQmqalsrJWOdQQrhc" + "L8TDZG4HBasd/kGr7EnlEbLKrGaWbAWYo4vssoUE7Q57wPumuMrYh4OKInksQTnB0OV9OIZEiTgb9fyaxJBCUiypuQDGMlAP8nb0TQF1vxjR" + "vsXV/qnN0ovR6/7raePN4NM4nMEXESRQyC22Ml88z+WZhe2VG5CLH8tFNapisH+msyqtGyDa+SEie/e6xcDfI1fLAZwEHM34/YjW+C1E1PPd" + "RsGBE9Ms6JUaPrcITwImVKMLGH1pj+TsS+qK7JTGjH8FKwR1Eds/dJTFDEQaAYcmjXWxDR5HBknx7xv90K2uoKt4KmAgulWCQv/LaEAt27lm" + "aGP+4PQ9tpb/Otsub8spmcXKfNqdjWTDH3aEF5MsJ5eSddn92pFV9JAo4oFcj1fUTujfwumGCYDADjLBmyvQao9S42zpO5idmYrHWHt7h7m6" + "BoU5b4uKVS4w4ELEl8y+Oq1uiQLSB8e0cxLarWfCiiTL6v6aJUDEvUTQU+O3XiOxNMCkc00J7e1TnNnnoFqs0sqhV0mMT346vLHocBMVn2iE" + "ZSwFb9OWlHO1e9Y7ccd8XL7KjXtoZeuBTrdYlxrf3g6QHVZg3OOM3OVSOWawFU0tBZUUgH6rYH4oHh8JqwGmbTI1mtAtUX89KrHzGWgUs619" + "SPbFptVh8Vfb3zr0D0OzDqjzv1KCNW/tZcJnCaL9Zl6tAt2v5FiBTOjIbBgB" +) + +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/chat/completions", + litellm_request={**CHAT_COMPLETIONS_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request=VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8AfAmbYAQCTSseeQSZJ", + "content": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually in my " + "available toolset — I only have search_files and read_file. I'll note that discrepancy and proceed by " + "explaining which file I need, then calling read_file on shop/validation.py, probably starting around the " + "first 60 lines to get context near line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` next, starting at the `validate_quantity` definition on line 18 and " + "including the surrounding context. The project instructions mention a chat_completions client, but none " + "is available in my tools, so I'm using the read-only file tools." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_017jjdVAoN8bus5ng5NkzafQ", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 876, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 876, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 348, + "output_tokens_details": {"thinking_tokens": 168}, + }, + }, + expected_litellm_response={ + **CHAT_COMPLETIONS_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": ( + "I'll read `shop/validation.py` next, starting at the `validate_quantity` definition on line 18 and " + "including the surrounding context. The project instructions mention a chat_completions client, but " + "none is available in my tools, so I'm using the read-only file tools." + ), + "role": "assistant", + "tool_calls": [ + { + "index": 2, + "function": { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 60}', + "name": "read_file", + }, + "id": "toolu_vrtx_017jjdVAoN8bus5ng5NkzafQ", + "type": "function", + } + ], + "reasoning_content": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually in my " + "available toolset — I only have search_files and read_file. I'll note that discrepancy and proceed " + "by explaining which file I need, then calling read_file on shop/validation.py, probably starting " + "around the first 60 lines to get context near line 18.\n\n" + ), + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually " + "in my available toolset — I only have search_files and read_file. I'll note that discrepancy " + "and proceed by explaining which file I need, then calling read_file on shop/validation.py, " + "probably starting around the first 60 lines to get context near line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + } + ], + "provider_specific_fields": { + "citations": None, + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not " + "actually in my available toolset — I only have search_files and read_file. I'll note " + "that discrepancy and proceed by explaining which file I need, then calling read_file on " + "shop/validation.py, probably starting around the first 60 lines to get context near " + "line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + } + ], + }, + }, + } + ], + "usage": { + "completion_tokens": 348, + "prompt_tokens": 880, + "total_tokens": 1228, + "completion_tokens_details": {"reasoning_tokens": 168, "text_tokens": 180}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 4, + "cache_write_tokens": 876, + "cache_creation_tokens": 876, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 876, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 876, + "cache_read_input_tokens": 0, + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/chat/completions", diff --git a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py index 71cbd84543e..7b7f5a7876c 100644 --- a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py +++ b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.chat_completions.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/integration/translation/chat_completions/cross_model/README.md b/tests/integration/translation/chat_completions/cross_model/README.md index 0adce683f38..ee7c178da81 100644 --- a/tests/integration/translation/chat_completions/cross_model/README.md +++ b/tests/integration/translation/chat_completions/cross_model/README.md @@ -5,10 +5,10 @@ Every row sends `CODING_AGENT_REQUEST` to `/v1/chat/completions`, changing only | Tier | Model versions | | --- | --- | | Small | Haiku 4.5 and 5.5, GPT 5 Mini and 5.4 Mini, GPT 5.6 Luna and 6 Luna | -| Mid | Opus 5 and 5.5, GPT 5.6 Sol, 6 Sol and 6.1 Sol, GPT 5.6 Terra | +| Mid | Opus 5 and 5.5, Sonnet 5.5, Gemini 3.8 Flash and 3.1 Pro Preview, GPT 5.6 Sol, 6 Sol and 6.1 Sol, GPT 5.6 Terra | | High | Fable 5 and 5.1, Gemini 4 Argon, GPT 6 Astra | -The 45 rows cover Anthropic direct, OpenAI direct, Vertex, Azure OpenAI, Azure AI and Bedrock Converse where accessible. Vertex Haiku 5.5 was unavailable during capture and is not represented +The 48 rows cover Anthropic direct, OpenAI direct, Vertex, Azure OpenAI, Azure AI and Bedrock Converse where accessible. Vertex Haiku 5.5 was unavailable during capture and is not represented The request contains system instructions, conversation history, previous assistant tool calls and results, four function definitions, automatic parallel tools, low reasoning effort and an output limit. Each row exercises a tool-call response, raw streaming, SDK stream assembly and a follow-up that replays the returned assistant message and tool-call IDs @@ -33,7 +33,7 @@ Streaming cases store `expected_litellm_message` and `expected_litellm_usage`. M Usage requires nonnegative integer counts, positive prompt and completion counts, and `total_tokens == prompt_tokens + completion_tokens`. The complete normalized usage must equal the independently expected counters, including cached and reasoning tokens and derived text/cache details. Missing, changed or unexpected counters fail. Only null optional fields, the explicitly listed provider metadata/duplicate counters and computed `cost` are excluded. Thus a wrong prompt count with an adjusted total still fails -The raw stream check assembles tool names and argument fragments by index. The SDK check separately uses `OpenAI.chat.completions.stream()` without removing reasoning fields. All 45 SDK checks run, including the seven Haiku rows. [LIT-9653](https://linear.app/litellm-ai/issue/LIT-9653) was fixed in [#46611](https://github.com/BerriAI/litellm/pull/46611), so there are no Haiku skips +The raw stream check assembles tool names and argument fragments by index. The SDK check separately uses `OpenAI.chat.completions.stream()` without removing reasoning fields. All 48 SDK checks run, including the seven Haiku rows. [LIT-9653](https://linear.app/litellm-ai/issue/LIT-9653) was fixed in [#46611](https://github.com/BerriAI/litellm/pull/46611), so there are no Haiku skips ## Recorded upstreams diff --git a/tests/integration/translation/chat_completions/cross_model/anthropic.py b/tests/integration/translation/chat_completions/cross_model/anthropic.py index ccb543ccf51..3df7da17ef6 100644 --- a/tests/integration/translation/chat_completions/cross_model/anthropic.py +++ b/tests/integration/translation/chat_completions/cross_model/anthropic.py @@ -2,10 +2,7 @@ from dataclasses import replace from typing import Final from integration._support.client import list_value, object_value -from integration.translation.case import ( - SSEFrame, - StreamingTranslationTestCase, -) +from integration.translation.case import SSEFrame, StreamingTranslationTestCase from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.bases.anthropic import ( CLAUDE_FABLE_5_1_TEST_CASE, @@ -20,10 +17,7 @@ from integration.translation.chat_completions.templates import ( coding_agent_stream_message, CODING_AGENT_RESPONSE, ) -from integration.translation.templates.anthropic import ( - ANTHROPIC_AGENT_REQUEST, - ANTHROPIC_AGENT_RESPONSE, -) +from integration.translation.templates.anthropic import ANTHROPIC_AGENT_REQUEST, ANTHROPIC_AGENT_RESPONSE CLAUDE_FABLE_5_1_AGENT_TEST_CASE: Final = replace( CLAUDE_FABLE_5_1_TEST_CASE, @@ -350,7 +344,7 @@ CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01324ezmdSmGSpmxGX55GvYF", "toolu_01MxHYehH7UYVY2A8WkRJZPu"), + ("toolu_01324ezmdSmGSpmxGX55GvYF", "toolu_01MxHYehH7UYVY2A8WkRJZPu") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -596,7 +590,7 @@ CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01UXBJZGdCaY4vxon1j4uapy", "toolu_0124RYbAPnwVoGtdzfPCNPc6"), + ("toolu_01UXBJZGdCaY4vxon1j4uapy", "toolu_0124RYbAPnwVoGtdzfPCNPc6") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -1494,7 +1488,7 @@ CLAUDE_HAIKU_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_0177HUb25ArysRAu8eF3u4DJ", "toolu_011gegGWvYoxPwxWY7LSK4Uk"), + ("toolu_0177HUb25ArysRAu8eF3u4DJ", "toolu_011gegGWvYoxPwxWY7LSK4Uk") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -1837,7 +1831,7 @@ CLAUDE_OPUS_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_013mPRCxrKKH6KzYGHKPukwv", "toolu_01FkXiaw4cTiXua3PoMd4wCu"), + ("toolu_013mPRCxrKKH6KzYGHKPukwv", "toolu_01FkXiaw4cTiXua3PoMd4wCu") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -2088,7 +2082,7 @@ CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_017r3PoQsgodP2mgcgDnDLxU", "toolu_01T8ECEb4XwFyozSnCKUiuB2"), + ("toolu_017r3PoQsgodP2mgcgDnDLxU", "toolu_01T8ECEb4XwFyozSnCKUiuB2") ), expected_litellm_usage={ "prompt_tokens": 1180, diff --git a/tests/integration/translation/chat_completions/cross_model/azure_ai.py b/tests/integration/translation/chat_completions/cross_model/azure_ai.py index 72c14ef103c..73d0ffbb8df 100644 --- a/tests/integration/translation/chat_completions/cross_model/azure_ai.py +++ b/tests/integration/translation/chat_completions/cross_model/azure_ai.py @@ -2,10 +2,7 @@ from dataclasses import replace from typing import Final from integration._support.client import list_value, object_value -from integration.translation.case import ( - SSEFrame, - StreamingTranslationTestCase, -) +from integration.translation.case import SSEFrame, StreamingTranslationTestCase from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.bases.azure_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, @@ -18,10 +15,7 @@ from integration.translation.chat_completions.templates import ( coding_agent_stream_message, CODING_AGENT_RESPONSE, ) -from integration.translation.templates.anthropic import ( - ANTHROPIC_AGENT_REQUEST, - ANTHROPIC_AGENT_RESPONSE, -) +from integration.translation.templates.anthropic import ANTHROPIC_AGENT_REQUEST, ANTHROPIC_AGENT_RESPONSE CLAUDE_FABLE_5_1_AGENT_TEST_CASE: Final = replace( CLAUDE_FABLE_5_1_TEST_CASE, @@ -349,7 +343,7 @@ CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01PfMSQLj6cQzebc2vGbdsV8", "toolu_01SVc4RwmEaMiDK33hUEbTZW"), + ("toolu_01PfMSQLj6cQzebc2vGbdsV8", "toolu_01SVc4RwmEaMiDK33hUEbTZW") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -593,7 +587,7 @@ CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01DprHn2GeJDAWFEyTqBzy2Z", "toolu_016kava9hqtKw71W4vBEedfW"), + ("toolu_01DprHn2GeJDAWFEyTqBzy2Z", "toolu_016kava9hqtKw71W4vBEedfW") ), expected_litellm_usage={ "prompt_tokens": 1184, diff --git a/tests/integration/translation/chat_completions/cross_model/case.py b/tests/integration/translation/chat_completions/cross_model/case.py index d0e842cfa67..d81e6a269f6 100644 --- a/tests/integration/translation/chat_completions/cross_model/case.py +++ b/tests/integration/translation/chat_completions/cross_model/case.py @@ -1,9 +1,6 @@ from typing import Final -from integration.translation.case import ( - StreamingTranslationTestCase, - TranslationTestCase, -) +from integration.translation.case import StreamingTranslationTestCase, TranslationTestCase from integration.translation.chat_completions.cross_model import ( anthropic, azure, @@ -180,6 +177,21 @@ MID: Final[tuple[AgentCases, ...]] = ( openai.GPT_6_1_SOL_AGENT_FOLLOWUP_TEST_CASE, openai.GPT_6_1_SOL_AGENT_STREAMING_TEST_CASE, ), + ( + vertex_ai.CLAUDE_SONNET_5_5_AGENT_TEST_CASE, + vertex_ai.CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE, + vertex_ai.CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE, + ), + ( + vertex_ai.GEMINI_3_8_FLASH_AGENT_TEST_CASE, + vertex_ai.GEMINI_3_8_FLASH_AGENT_FOLLOWUP_TEST_CASE, + vertex_ai.GEMINI_3_8_FLASH_AGENT_STREAMING_TEST_CASE, + ), + ( + vertex_ai.GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE, + vertex_ai.GEMINI_3_1_PRO_PREVIEW_AGENT_FOLLOWUP_TEST_CASE, + vertex_ai.GEMINI_3_1_PRO_PREVIEW_AGENT_STREAMING_TEST_CASE, + ), ) HIGH: Final[tuple[AgentCases, ...]] = ( diff --git a/tests/integration/translation/chat_completions/cross_model/contracts.py b/tests/integration/translation/chat_completions/cross_model/contracts.py index 4fbf6a880b5..a15a43985f1 100644 --- a/tests/integration/translation/chat_completions/cross_model/contracts.py +++ b/tests/integration/translation/chat_completions/cross_model/contracts.py @@ -121,7 +121,9 @@ PROVIDER_EXTENSIONS: Final = MappingProxyType( ("provider_specific_fields", "reasoning_content", "thinking_blocks", "reasoning_items", "images") ), "choices[].message.tool_calls[]": frozenset(("index", "caller", "provider_specific_fields")), - "usage": frozenset(("cache_creation_input_tokens", "cache_read_input_tokens", "inference_geo", "service_tier")), + "usage": frozenset( + ("cache_creation_input_tokens", "cache_read_input_tokens", "inference_geo", "service_tier", "speed") + ), } ) diff --git a/tests/integration/translation/chat_completions/cross_model/vertex_ai.py b/tests/integration/translation/chat_completions/cross_model/vertex_ai.py index 42f363b1920..0dbb4b277c5 100644 --- a/tests/integration/translation/chat_completions/cross_model/vertex_ai.py +++ b/tests/integration/translation/chat_completions/cross_model/vertex_ai.py @@ -6,19 +6,22 @@ from integration.translation.case import ( SSEFrame, StreamingTranslationTestCase, ) -from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_FABLE_5_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_OPUS_5_TEST_CASE, + CLAUDE_SONNET_5_5_TEST_CASE, + GEMINI_3_1_PRO_PREVIEW_TEST_CASE, + GEMINI_3_8_FLASH_TEST_CASE, GEMINI_4_ARGON_TEST_CASE, ) +from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.templates import ( CODING_AGENT_REQUEST, - coding_agent_stream_message, CODING_AGENT_RESPONSE, + coding_agent_stream_message, ) from integration.translation.templates.anthropic import ( ANTHROPIC_AGENT_REQUEST, @@ -1721,6 +1724,1214 @@ CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( ) +CLAUDE_SONNET_5_5_AGENT_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "vertex_ai/claude-sonnet-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-sonnet-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + "anthropic_version": "vertex-2023-10-16", + }, + mock_provider_response={ + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_vrtx_011Cg63qBoZy88hTvxUdtSEr", + "content": [ + { + "type": "tool_use", + "id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "name": "search_code", + "input": {"query": "retries_remaining"}, + }, + { + "type": "tool_use", + "id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "name": "run_tests", + "input": {"path": "tests/test_retry.py"}, + }, + ], + "stop_reason": "tool_use", + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "model": "vertex_ai/claude-sonnet-5-5@default", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "function": {"arguments": '{"query": "retries_remaining"}', "name": "search_code"}, + "id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "type": "function", + }, + { + "index": 1, + "function": {"arguments": '{"path": "tests/test_retry.py"}', "name": "run_tests"}, + "id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "type": "function", + }, + ], + "provider_specific_fields": {"citations": None, "thinking_blocks": None}, + "content": None, + }, + } + ], + "usage": { + "completion_tokens": 105, + "prompt_tokens": 1184, + "total_tokens": 1289, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + }, +) + + +CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_request, CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_litellm_response + ), + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-sonnet-5-5@default:rawPredict", + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "messages": [ + *list_value(ANTHROPIC_AGENT_REQUEST["messages"]), + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "name": "search_code", + "input": {"query": "retries_remaining"}, + }, + { + "type": "tool_use", + "id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "name": "run_tests", + "input": {"path": "tests/test_retry.py"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "content": '{"matches": [{"path": "src/worker.py", "line": 12, "text": "remaining = ' + 'retries_remaining(3, attempts)"}]}', + }, + { + "type": "tool_result", + "tool_use_id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "content": '{"exit_code": 1, "output": "FAILED tests/test_retry.py::test_retries_remaining - assert ' + '4 == 2"}', + }, + { + "type": "text", + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences.", + }, + ], + }, + ], + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + "anthropic_version": "vertex-2023-10-16", + }, + mock_provider_response={ + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_vrtx_011Cg63wSU3MbRaa6kwyvRLk", + "content": [ + { + "type": "text", + "text": "The bug is that `retries_remaining` in `src/retry.py` returns `limit + attempts` " + "instead of `max(0, limit - attempts)`, so it overcounts and can't be clamped at zero. " + "The one other caller is `src/worker.py` line 12 (`retries_remaining(3, attempts)`), and " + "`tests/test_retry.py::test_retries_remaining` fails with `assert 4 == 2`.", + } + ], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1492, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 138, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "model": "vertex_ai/claude-sonnet-5-5@default", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The bug is that `retries_remaining` in `src/retry.py` returns `limit + attempts` " + "instead of `max(0, limit - attempts)`, so it overcounts and can't be clamped at zero. " + "The one other caller is `src/worker.py` line 12 (`retries_remaining(3, attempts)`), and " + "`tests/test_retry.py::test_retries_remaining` fails with `assert 4 == 2`.", + "role": "assistant", + "provider_specific_fields": {"citations": None, "thinking_blocks": None}, + }, + } + ], + "usage": { + "completion_tokens": 138, + "prompt_tokens": 1492, + "total_tokens": 1630, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 138}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1492, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + }, +) + +CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-sonnet-5-5@default:streamRawPredict?alt=sse", + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + "anthropic_version": "vertex-2023-10-16", + "stream": True, + }, + mock_provider_frames=( + SSEFrame( + data={ + "type": "message_start", + "message": { + "model": "claude-sonnet-5-5", + "id": "msg_vrtx_011Cg63qMXiZfvJr7NuvT9B1", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 16, + }, + }, + }, + event="message_start", + ), + SSEFrame(data={"type": "ping"}, event="ping"), + SSEFrame( + data={ + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_vrtx_01EJMuteinVLQX3UQ8bxJEps", + "name": "search_code", + "input": {}, + }, + }, + event="content_block_start", + ), + SSEFrame( + data={"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '": "ret'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": "ries_re"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": "maining"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '"}'}, + }, + event="content_block_delta", + ), + SSEFrame(data={"type": "content_block_stop", "index": 0}, event="content_block_stop"), + SSEFrame( + data={ + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_vrtx_015Mb3cdfYTHCA9QowGSX2RP", + "name": "run_tests", + "input": {}, + }, + }, + event="content_block_start", + ), + SSEFrame( + data={"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": ""}}, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '{"path"'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": ': "test'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": "s/test_"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": "retry.p"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": 'y"}'}, + }, + event="content_block_delta", + ), + SSEFrame(data={"type": "content_block_stop", "index": 1}, event="content_block_stop"), + SSEFrame( + data={ + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + event="message_delta", + ), + SSEFrame(data={"type": "message_stop"}, event="message_stop"), + ), + expected_litellm_message=coding_agent_stream_message( + ("toolu_vrtx_01EJMuteinVLQX3UQ8bxJEps", "toolu_vrtx_015Mb3cdfYTHCA9QowGSX2RP") + ), + expected_litellm_usage={ + "prompt_tokens": 1184, + "completion_tokens": 105, + "total_tokens": 1289, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + }, +) + + +GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE: Final = replace( + GEMINI_3_1_PRO_PREVIEW_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "vertex_ai/gemini-3.1-pro-preview"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent", + expected_provider_headers={"content-type": "application/json", "authorization": "Bearer scripted-token"}, + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + "thought": True, + }, + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_2861166", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_2861167", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 551, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + "thoughtsTokenCount": 106, + }, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:08.649083Z", + "responseId": "QDDSavvOJ6-q_9MP5byA4Qo", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "QDDSavvOJ6-q_9MP5byA4Qo", + "model": "vertex_ai/gemini-3.1-pro-preview", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "provider_specific_fields": {"thought_signature": "synthetic-signature"}, + "function": {"arguments": '{"path": "tests/test_retry.py"}', "name": "run_tests"}, + "id": "call_2861166__thought__synthetic-signature", + "type": "function", + }, + { + "index": 1, + "function": {"arguments": '{"query": "retries_remaining"}', "name": "search_code"}, + "id": "call_2861167", + "type": "function", + }, + ], + "images": [], + "reasoning_content": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + } + ], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + "content": None, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 147, + "prompt_tokens": 404, + "total_tokens": 551, + "completion_tokens_details": {"reasoning_tokens": 106, "text_tokens": 41}, + "prompt_tokens_details": {"text_tokens": 404}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + + +GEMINI_3_1_PRO_PREVIEW_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.litellm_request, + GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.expected_litellm_response, + ), + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent", + expected_provider_request={ + **VERTEX_AI_AGENT_REQUEST, + "contents": [ + *list_value(VERTEX_AI_AGENT_REQUEST["contents"]), + { + "role": "model", + "parts": [ + { + "thought": True, + "text": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + }, + { + "function_call": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_2861166", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "function_call": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_2861167", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "name": "run_tests", + "response": { + "exit_code": 1, + "output": "FAILED tests/test_retry.py::test_retries_remaining - assert 4 == 2", + }, + "id": "call_2861166", + } + }, + { + "function_response": { + "name": "search_code", + "response": { + "matches": [ + { + "path": "src/worker.py", + "line": 12, + "text": "remaining = retries_remaining(3, attempts)", + } + ] + }, + "id": "call_2861167", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences." + } + ], + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "**My Debugging Session**\n" + "\n" + "It appears I've stumbled upon a rather straightforward arithmetic error within the " + "`retries_remaining` function. My current understanding is that it's incorrectly adding " + "the `limit` and `attempts` parameters instead of performing a subtraction, and " + "consequently, it's not ensuring the result is capped at zero. I've identified that " + "`src/worker.py` is a caller of this function, and specifically, the " + "`tests/test_retry.py::test_retries_remaining` test is failing because, for instance, " + "`retries_remaining(3, 1)` is returning 4 when it should be returning 2.\n" + "\n", + "thought": True, + }, + { + "text": "The bug in `retries_remaining` is that it adds `limit` and `attempts` instead of " + "calculating `max(0, limit - attempts)`. This causes `tests/test_retry.py` to fail, " + "which was verified, and the function is called by `src/worker.py`.", + "thoughtSignature": "synthetic-signature", + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 525, + "candidatesTokenCount": 64, + "totalTokenCount": 706, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 525}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 64}], + "thoughtsTokenCount": 117, + }, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:11:28.565345Z", + "responseId": "kDDSauHAIq6d4_UPv9qnMQ", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "kDDSauHAIq6d4_UPv9qnMQ", + "model": "vertex_ai/gemini-3.1-pro-preview", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The bug in `retries_remaining` is that it adds `limit` and `attempts` instead of " + "calculating `max(0, limit - attempts)`. This causes `tests/test_retry.py` to fail, " + "which was verified, and the function is called by `src/worker.py`.", + "role": "assistant", + "images": [], + "reasoning_content": "**My Debugging Session**\n" + "\n" + "It appears I've stumbled upon a rather straightforward arithmetic error within the " + "`retries_remaining` function. My current understanding is that it's incorrectly adding " + "the `limit` and `attempts` parameters instead of performing a subtraction, and " + "consequently, it's not ensuring the result is capped at zero. I've identified that " + "`src/worker.py` is a caller of this function, and specifically, the " + "`tests/test_retry.py::test_retries_remaining` test is failing because, for instance, " + "`retries_remaining(3, 1)` is returning 4 when it should be returning 2.\n" + "\n", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "**My Debugging Session**\n" + "\n" + "It appears I've stumbled upon a rather straightforward arithmetic error within the " + "`retries_remaining` function. My current understanding is that it's incorrectly adding " + "the `limit` and `attempts` parameters instead of performing a subtraction, and " + "consequently, it's not ensuring the result is capped at zero. I've identified that " + "`src/worker.py` is a caller of this function, and specifically, the " + "`tests/test_retry.py::test_retries_remaining` test is failing because, for instance, " + "`retries_remaining(3, 1)` is returning 4 when it should be returning 2.\n" + "\n", + } + ], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 181, + "prompt_tokens": 525, + "total_tokens": 706, + "completion_tokens_details": {"reasoning_tokens": 117, "text_tokens": 64}, + "prompt_tokens_details": {"text_tokens": 525}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + +GEMINI_3_1_PRO_PREVIEW_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:streamGenerateContent?alt=sse", + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_frames=( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_1016270", + }, + "thoughtSignature": "synthetic-signature", + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:13.305463Z", + "responseId": "RTDSarfSEuXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_1016276", + } + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:13.305463Z", + "responseId": "RTDSarfSEuXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 473, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + "thoughtsTokenCount": 28, + }, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:13.305463Z", + "responseId": "RTDSarfSEuXG4_UP_a-tuQo", + } + ), + ), + expected_litellm_message={ + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1016270__thought__synthetic-signature", + "type": "function", + "function": {"name": "run_tests", "arguments": '{"path": "tests/test_retry.py"}'}, + }, + { + "id": "call_1016276", + "type": "function", + "function": {"name": "search_code", "arguments": '{"query": "retries_remaining"}'}, + }, + ], + "reasoning_content": None, + "thinking_blocks": [], + "reasoning_items": [], + }, + expected_litellm_usage={ + "prompt_tokens": 404, + "completion_tokens": 69, + "total_tokens": 473, + "prompt_tokens_details": {"text_tokens": 404}, + "completion_tokens_details": {"reasoning_tokens": 28, "text_tokens": 41}, + }, +) + + +GEMINI_3_8_FLASH_AGENT_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "vertex_ai/gemini-3.8-flash"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.8-flash:generateContent", + expected_provider_headers={"content-type": "application/json", "authorization": "Bearer scripted-token"}, + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_1969624", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_1969625", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 445, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:02.505616Z", + "responseId": "OjDSapDuHpCG9LsPreTxQA", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "OjDSapDuHpCG9LsPreTxQA", + "model": "vertex_ai/gemini-3.8-flash", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "provider_specific_fields": {"thought_signature": "synthetic-signature"}, + "function": {"arguments": '{"query": "retries_remaining"}', "name": "search_code"}, + "id": "call_1969624__thought__synthetic-signature", + "type": "function", + }, + { + "index": 1, + "function": {"arguments": '{"path": "tests/test_retry.py"}', "name": "run_tests"}, + "id": "call_1969625", + "type": "function", + }, + ], + "images": [], + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + "content": None, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 41, + "prompt_tokens": 404, + "total_tokens": 445, + "completion_tokens_details": {"text_tokens": 41}, + "prompt_tokens_details": {"text_tokens": 404}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + + +GEMINI_3_8_FLASH_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + GEMINI_3_8_FLASH_AGENT_TEST_CASE.litellm_request, GEMINI_3_8_FLASH_AGENT_TEST_CASE.expected_litellm_response + ), + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.8-flash:generateContent", + expected_provider_request={ + **VERTEX_AI_AGENT_REQUEST, + "contents": [ + *list_value(VERTEX_AI_AGENT_REQUEST["contents"]), + { + "role": "model", + "parts": [ + { + "function_call": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_1969624", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "function_call": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_1969625", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "name": "search_code", + "response": { + "matches": [ + { + "path": "src/worker.py", + "line": 12, + "text": "remaining = retries_remaining(3, attempts)", + } + ] + }, + "id": "call_1969624", + } + }, + { + "function_response": { + "name": "run_tests", + "response": { + "exit_code": 1, + "output": "FAILED tests/test_retry.py::test_retries_remaining - assert 4 == 2", + }, + "id": "call_1969625", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences." + } + ], + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "`retries_remaining` erroneously adds `limit` and `attempts` instead of calculating " + "`max(0, limit - attempts)`, causing `assert retries_remaining(3, 1) == 2` to fail with " + "`assert 4 == 2`. Aside from its unit test, `src/worker.py` (line 12) was found as the " + "only other caller in the repository.", + "thoughtSignature": "synthetic-signature", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 525, + "candidatesTokenCount": 85, + "totalTokenCount": 610, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 525}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 85}], + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:11:25.313560Z", + "responseId": "jTDSatiRE9ap4_UPuPC7wQo", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "jTDSatiRE9ap4_UPuPC7wQo", + "model": "vertex_ai/gemini-3.8-flash", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "`retries_remaining` erroneously adds `limit` and `attempts` instead of calculating " + "`max(0, limit - attempts)`, causing `assert retries_remaining(3, 1) == 2` to fail with " + "`assert 4 == 2`. Aside from its unit test, `src/worker.py` (line 12) was found as the " + "only other caller in the repository.", + "role": "assistant", + "images": [], + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 85, + "prompt_tokens": 525, + "total_tokens": 610, + "completion_tokens_details": {"text_tokens": 85}, + "prompt_tokens_details": {"text_tokens": 525}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + +GEMINI_3_8_FLASH_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=GEMINI_3_8_FLASH_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **GEMINI_3_8_FLASH_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=GEMINI_3_8_FLASH_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_frames=( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_302633", + }, + "thoughtSignature": "synthetic-signature", + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:05.862754Z", + "responseId": "PTDSaqLUNOXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_302642", + } + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:05.862754Z", + "responseId": "PTDSaqLUNOXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 445, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:05.862754Z", + "responseId": "PTDSaqLUNOXG4_UP_a-tuQo", + } + ), + ), + expected_litellm_message=coding_agent_stream_message(("call_302633__thought__synthetic-signature", "call_302642")), + expected_litellm_usage={ + "prompt_tokens": 404, + "completion_tokens": 41, + "total_tokens": 445, + "prompt_tokens_details": {"text_tokens": 404}, + "completion_tokens_details": {"text_tokens": 41}, + }, +) + + GEMINI_4_ARGON_AGENT_TEST_CASE: Final = replace( GEMINI_4_ARGON_TEST_CASE, scenario="agent_tools", diff --git a/tests/integration/translation/messages/bases/anthropic.py b/tests/integration/translation/messages/bases/anthropic.py index cc542499a07..be7f487f1b0 100644 --- a/tests/integration/translation/messages/bases/anthropic.py +++ b/tests/integration/translation/messages/bases/anthropic.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from typing import Final +from integration._support.client import object_value from integration.translation.case import TranslationTestCase from integration.translation.messages.templates import MESSAGES_REQUEST, MESSAGES_RESPONSE from integration.translation.templates.anthropic import ANTHROPIC_MESSAGES_REQUEST, ANTHROPIC_MESSAGES_RESPONSE @@ -31,7 +32,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CffzUNHaEfzVxCh5hskBG", + "id": "msg_011Cg7v1w1xZexMrWpf5GU75", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -47,12 +48,13 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-sonnet-4-6", - "id": "msg_011CffzUNHaEfzVxCh5hskBG", + "id": "msg_011Cg7v1w1xZexMrWpf5GU75", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -68,6 +70,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -98,10 +101,10 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011CfgDz9Vj8wXvKzhoM4MQA", + "id": "msg_011Cg7v1Mt9Eid844XTV8AMc", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -114,15 +117,16 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-haiku-4-5", - "id": "msg_011CfgDz9Vj8wXvKzhoM4MQA", + "id": "msg_011Cg7v1Mt9Eid844XTV8AMc", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -135,6 +139,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -165,7 +170,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-5", - "id": "msg_011CfgDzLXmJmnHRTpeBoDmZ", + "id": "msg_011Cg7v2KRMCpM2FsiVygFAS", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -182,12 +187,13 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-sonnet-5", - "id": "msg_011CfgDzLXmJmnHRTpeBoDmZ", + "id": "msg_011Cg7v2KRMCpM2FsiVygFAS", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -204,6 +210,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -234,7 +241,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfgDzSFYYPC3aYE9rH98J", + "id": "msg_011Cg7v1ZPwvewtTtXqjdwzZ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -251,12 +258,13 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-opus-4-8", - "id": "msg_011CfgDzSFYYPC3aYE9rH98J", + "id": "msg_011Cg7v1ZPwvewtTtXqjdwzZ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -273,6 +281,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -377,11 +386,13 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_SONNET_5_5_RESPONSE_FIELDS, "model": "claude-sonnet-5-5", + "usage": {**object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_SONNET_5_5_RESPONSE_FIELDS, "model": "anthropic/claude-sonnet-5-5", + "usage": {**object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -460,11 +471,13 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "anthropic/claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -555,11 +568,13 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_OPUS_5_5_RESPONSE_FIELDS, "model": "claude-opus-5-5", + "usage": {**object_value(CLAUDE_OPUS_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_OPUS_5_5_RESPONSE_FIELDS, "model": "anthropic/claude-opus-5-5", + "usage": {**object_value(CLAUDE_OPUS_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -651,10 +666,12 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "anthropic/claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) diff --git a/tests/integration/translation/messages/bases/azure_ai.py b/tests/integration/translation/messages/bases/azure_ai.py index 454a76400db..6de15226c48 100644 --- a/tests/integration/translation/messages/bases/azure_ai.py +++ b/tests/integration/translation/messages/bases/azure_ai.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from typing import Final +from integration._support.client import object_value from integration.translation.case import TranslationTestCase from integration.translation.messages.templates import MESSAGES_REQUEST, MESSAGES_RESPONSE from integration.translation.templates.anthropic import ANTHROPIC_MESSAGES_REQUEST, ANTHROPIC_MESSAGES_RESPONSE @@ -31,10 +32,10 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011Cfjax2zxeHkddgerkEGRY", + "id": "msg_011Cg7vB62SdiZVj3M7xJXpr", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -47,15 +48,16 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "azure_ai/claude-haiku-4-5", - "id": "msg_011Cfjax2zxeHkddgerkEGRY", + "id": "msg_011Cg7vB62SdiZVj3M7xJXpr", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -68,6 +70,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -98,7 +101,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfjaxAknNm3ZLkVxQRj8M", + "id": "msg_011Cg7vCCFGaPEhwGwoS6VPy", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -114,12 +117,13 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "azure_ai/claude-sonnet-4-6", - "id": "msg_011CfjaxAknNm3ZLkVxQRj8M", + "id": "msg_011Cg7vCCFGaPEhwGwoS6VPy", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -135,6 +139,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -234,7 +239,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfjaxPE8FBKSJ9KbtQzDE", + "id": "msg_011Cg7vBVpRK6kd1kkYXJXyb", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -251,12 +256,13 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "azure_ai/claude-opus-4-8", - "id": "msg_011CfjaxPE8FBKSJ9KbtQzDE", + "id": "msg_011Cg7vBVpRK6kd1kkYXJXyb", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -273,6 +279,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -483,11 +490,13 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "azure_ai/claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -541,10 +550,12 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "azure_ai/claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) diff --git a/tests/integration/translation/messages/bases/vertex_ai.py b/tests/integration/translation/messages/bases/vertex_ai.py index 6f506b39d93..120b764a63c 100644 --- a/tests/integration/translation/messages/bases/vertex_ai.py +++ b/tests/integration/translation/messages/bases/vertex_ai.py @@ -10,6 +10,7 @@ from integration.translation.templates.vertex_ai import ( ) GEMINI_3_5_FLASH_THOUGHT_SIGNATURE: Final = "AY89a1+kriNs9YsPegJQxgjTWsMlxsqPU/UkHDyz4EK+3FRrE52ygqPIHTvkblgKKiO5Fn8R9D3iu//i+YnuuZ+e7W0yxUoFScas89b5rn3vZHsihnUA7VTe3b7beth1IftteoTGZ3F8OSLrSIP8HKn4xTIx35eeNEt5oiD3K5rTQrdklqO277HvTDhu/LuZhZ0pzaWciGQEc61epIBBpC2ygu8V4L8TekUVZnz6qzSzxipdiQStQ8dZou43334hMnxOR6JjJCGxqt/1K+C0vXX6YP6f+Bph5VhPSnvy2Z1uApi3TUxAcsDDyVewIr+ha2G6DPhN1Oi5uZZNhmbvH01Uq5OOquF1uEICq2m/35LhJIbDNM9N7oI1Z6L94P9xdugeFXM2S7wdJfc8sdg09vEzNeA++GK25zRuB4hU7XyaG4XnWjmu2kswi4sf+EyrO/K0P2tb/ZCJu7kJTtWTiD1wgSIl1eqS8ORKKsXDN2hIbFrtsNR50WIvVzd6" +CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE: Final = "CAQS4wYKEggSEAIYAjgBQgh0aGlua2luZxIMgzBnOw4aYVVqPKRPGgyphQhlJ0pt7tgasWIiMOWM4ATud/TWad7qo/Eh9ZSXhK3IcgnwdgqVeXog9D1q/y69qnyrHdN3a8RYUrAgtyr+BYkQvRW/cF0vH0l86vIM95N77usMXdZKRX3I8XW68KpQedYwFoR12sS9q7A24CcN5UMyQljcF76GhmAFnvpk4yEs7pk8b7ZRSXzBsU6wJfMYCoGXC9r3hZN8m4sBHLJGJRv5ekzEVGVvGCYooE1fBZ8ik/j2DCNOyVdDs4/7xzkqRo7dgFsNJI9Nm7+HgwEdvb1rORObJWooeYb8tfofTFVO43qvNdvM2p2d1Dcf6lsQmvE5fgDFc0t75AuvSZUrgDNe2bMXd2USxxL9YtJ1I9ae96hZAnyOqYZlyiuZf5lZtQQSRouVIdZvFUNonOHLy/Uyw/BGF2bO1/5rIFJBDJC8C6BpbyfDQPpal4YkALulDdgeJWI6dmJYUiK+A+3zcVgQh7at5JPbG5n0tQgO2IFkayqLjCf67AaGS6Xq8J97F+FSwppKvS0TvRbTEvIK1iO3HZZRwcRIcFBn5Llvv6R9JFknXEtohUqvfTnZfuPQ104vxqm91G2rFK1R6C1md6Ze3pJWD0wgkgZyQ7PIKsAmBF1qNoniFi5GVRJx6XHSzr4ZAPDDO1aYABIRl5OWilrkQP3xedj5pcAjJqFuto3cXlhQSNUrocOpEV5IIXxSJc5g5VwYAybpMU4FVImzkCJGJaHLWQKRo0k4IcvXLxBPNXZhop2CkEeazZskF4de0jn1dJa3sQ9rT+7fNVEWPvP+lkBqu/lfObBPVh5jGeouDrRcJvyeYKNZ0NQBdznM5rv4rN8HczeYg1WaYuNOMBvkqT5Zg/UhLHhHV7NWBSa0zt9DPevU/S3OBR1E0YIzCtPFTa0Ia4gGTX40hEc+supCA3h4HpehqMqETGe++Ah+U10HRhlb6lVfRhkU0ypdPqtpodgw/OdUOy7f3/sc31xknudQMoEQqv48F4gR34ee18Op+w1MGn0FAzAFE95D2BuJObRqUhZ6N6nWtlWOqJjCPR+0OrwR3JlxWh19Vxps/ASL6HaJKjClVG2ffojvmVnY8aHjEfI8ifbHE6EYAQ==" CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", @@ -486,6 +487,106 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/messages", + litellm_request={**MESSAGES_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={"authorization": "Bearer scripted-token", "content-type": "application/json"}, + expected_provider_request={ + **VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + "messages": MESSAGES_REQUEST["messages"], + "stream": False, + "system": [ + { + "type": "text", + "text": ( + "You help maintain a synthetic Python project. Read the relevant files before suggesting a fix. " + "Explain which file you need next, then call a read-only tool. Never change files. Use the messages " + "client for this project review. Consider zero, negative and non-integer quantities before choosing " + "the next file to inspect." + ), + }, + ], + "tools": MESSAGES_REQUEST["tools"], + }, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8BfjWDUi4ZYQxpy6aVM", + "content": [ + { + "type": "thinking", + "thinking": ( + "I found validate_quantity defined around line 18, so I'll read the validation.py file from the start " + "through roughly line 60 to capture the full context, keeping in mind I'll need to check how it handles " + "zero, negative, and non-integer cases.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` first. The search put `validate_quantity` at line 18, so I'll start " + "there and include the lines around it to see how zero, negative, and non-integer values are handled." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01YSY1kAV4ayCLVdBJkSv4dt", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 292, + "output_tokens_details": {"thinking_tokens": 129}, + }, + }, + expected_litellm_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "id": "msg_vrtx_011Cg8BfjWDUi4ZYQxpy6aVM", + "content": [ + { + "type": "thinking", + "thinking": ( + "I found validate_quantity defined around line 18, so I'll read the validation.py file from the start " + "through roughly line 60 to capture the full context, keeping in mind I'll need to check how it handles " + "zero, negative, and non-integer cases.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` first. The search put `validate_quantity` at line 18, so I'll start " + "there and include the lines around it to see how zero, negative, and non-integer values are handled." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01YSY1kAV4ayCLVdBJkSv4dt", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 292, + "output_tokens_details": {"thinking_tokens": 129}, + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", diff --git a/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py b/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py index a5b01dfd2e4..4f74a648d8e 100644 --- a/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py +++ b/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py index 5a3a760ffdc..6d7959b5de5 100644 --- a/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py @@ -8,6 +8,19 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.anthropic import CLAUDE_SONNET_4_6_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE from integration.translation.runner import assert_translation +CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS: Final = { + "content": [{"type": "text", "text": "Hello!"}], + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 5, + "service_tier": "standard", + "inference_geo": "global", + }, +} + CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( CLAUDE_SONNET_4_6_TEST_CASE, scenario="service_tier", @@ -16,9 +29,14 @@ CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request, "service_tier": "standard_only", }, - mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1RG36pCcTLE26htFFqC"}, + mock_provider_response={ + **CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, + "id": "msg_011Cg1RG36pCcTLE26htFFqC", + }, expected_litellm_response={ **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, "id": "msg_011Cg1RG36pCcTLE26htFFqC", }, ) diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py index 13bf9dd888f..f987f9d1f2a 100644 --- a/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py @@ -8,6 +8,19 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.azure_ai import CLAUDE_SONNET_4_6_TEST_CASE from integration.translation.runner import assert_translation +CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS: Final = { + "content": [{"type": "text", "text": "Hello!"}], + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 5, + "service_tier": "standard", + "inference_geo": "not_available", + }, +} + CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( CLAUDE_SONNET_4_6_TEST_CASE, scenario="service_tier", @@ -16,9 +29,14 @@ CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request, "service_tier": "standard_only", }, - mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1SnrstCdzZPwmTk2qQa"}, + mock_provider_response={ + **CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, + "id": "msg_011Cg1SnrstCdzZPwmTk2qQa", + }, expected_litellm_response={ **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, "id": "msg_011Cg1SnrstCdzZPwmTk2qQa", }, ) diff --git a/tests/integration/translation/response_contracts.py b/tests/integration/translation/response_contracts.py index c3182d269ff..6af3f9d681d 100644 --- a/tests/integration/translation/response_contracts.py +++ b/tests/integration/translation/response_contracts.py @@ -197,6 +197,7 @@ MESSAGES_OPTIONAL: Final = MappingProxyType( "server_tool_use", "service_tier", "inference_geo", + "speed", ) ), } diff --git a/tests/integration/translation/responses/bases/anthropic.py b/tests/integration/translation/responses/bases/anthropic.py index e30c1f15699..8d5051acb92 100644 --- a/tests/integration/translation/responses/bases/anthropic.py +++ b/tests/integration/translation/responses/bases/anthropic.py @@ -35,7 +35,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011CfgDzBnuPYc7TdysZFRpo", + "id": "msg_011Cg7v1CTrRytLeZfaZojBR", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -51,6 +51,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -130,10 +131,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfgDzH3AcaSNXtvurLgAG", + "id": "msg_011Cg7v1kr1sYXyQA143nsda", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello!"}], + "content": [{"type": "text", "text": "Hello! How can I help you?"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -143,9 +144,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - "output_tokens": 5, + "output_tokens": 11, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -164,7 +166,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "id": ANY, "status": "completed", "role": "assistant", - "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], + "content": [{"type": "output_text", "text": "Hello! How can I help you?", "annotations": []}], "phase": None, } ], @@ -190,9 +192,9 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "video_tokens": None, "cache_write_tokens": 0, }, - "output_tokens": 5, - "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 0, "text_tokens": 5}, - "total_tokens": 23, + "output_tokens": 11, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 0, "text_tokens": 11}, + "total_tokens": 29, "cost": None, }, "user": None, @@ -225,7 +227,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-5", - "id": "msg_011CfgDzPkj2HxBq8qGJsGSi", + "id": "msg_011Cg7v215mkRTDBRoPcerLf", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -242,6 +244,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -321,7 +324,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfgDzV9vZCkuQJt15n3Z1", + "id": "msg_011Cg7v1PgnYbGrw8eGpka6f", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -338,6 +341,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -462,6 +466,7 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 313}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={**RESPONSES_RESPONSE, "model": "anthropic/claude-sonnet-5-5"}, @@ -534,6 +539,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 67}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -553,7 +559,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "requested.\n" "\n", "annotations": [], - }, + } ], "phase": None, "encrypted_content": json.dumps( @@ -565,7 +571,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "requested.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], separators=(",", ":"), ), @@ -605,7 +611,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "range near the top of the file since no specific lines were requested.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, @@ -698,6 +704,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 130}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -719,7 +726,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "zero, negative, and non-integer inputs before moving to the test file.\n" "\n", "annotations": [], - }, + } ], "phase": None, "encrypted_content": json.dumps( @@ -842,6 +849,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -860,7 +868,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "rejection, then I'll read the test file in a separate step since there's " "nothing to derive from it for the read.", "annotations": [], - }, + } ], "phase": None, }, diff --git a/tests/integration/translation/responses/bases/azure_ai.py b/tests/integration/translation/responses/bases/azure_ai.py index 921fed14b45..dfcb454c482 100644 --- a/tests/integration/translation/responses/bases/azure_ai.py +++ b/tests/integration/translation/responses/bases/azure_ai.py @@ -31,7 +31,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011Cfjax6mBrNVhFGQopmVCp", + "id": "msg_011Cg7vB393hLWW1Xxq6d6vm", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -47,6 +47,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -127,7 +128,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfjaxGKedCbjtb7HuAUBC", + "id": "msg_011Cg7vCAx8gxwjwBr2R1kxz", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -143,6 +144,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -320,7 +322,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfjaxUiYCUsqMC4iqfc2U", + "id": "msg_011Cg7vBUGeH2HFZFTDxVdBh", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -337,6 +339,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -620,7 +623,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "suggesting a fix. Explain which file you need next, then call a read-only tool. " "Never change files. Use the responses client for this project review. Consider zero, " "negative and non-integer quantities before choosing the next file to inspect.", - }, + } ], }, mock_provider_response={ @@ -656,6 +659,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 73}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -673,7 +677,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "text": "I need to look at shop/validation.py to understand its contents, so I'll read " "through the first chunk of the file.\n\n", "annotations": [], - }, + } ], "phase": None, "encrypted_content": json.dumps( @@ -683,7 +687,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "thinking": "I need to look at shop/validation.py to understand its contents, so I'll read " "through the first chunk of the file.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], separators=(",", ":"), ), @@ -699,7 +703,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "text": "I'll read `shop/validation.py`, starting with the `validate_quantity` function found " "in the search.", "annotations": [], - }, + } ], "phase": None, }, @@ -737,7 +741,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "thinking": "I need to look at shop/validation.py to understand its contents, so I'll read " "through the first chunk of the file.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, @@ -764,7 +768,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "suggesting a fix. Explain which file you need next, then call a read-only tool. " "Never change files. Use the responses client for this project review. Consider zero, " "negative and non-integer quantities before choosing the next file to inspect.", - }, + } ], }, mock_provider_response={ @@ -801,6 +805,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, expected_litellm_response={ @@ -818,7 +823,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "text": "I'll read the validation function; reading the test file at the same time is " "independent of it, so I'll fetch both now.", "annotations": [], - }, + } ], "phase": None, }, diff --git a/tests/integration/translation/responses/bases/vertex_ai.py b/tests/integration/translation/responses/bases/vertex_ai.py index 734ebf13613..59637e380ce 100644 --- a/tests/integration/translation/responses/bases/vertex_ai.py +++ b/tests/integration/translation/responses/bases/vertex_ai.py @@ -12,6 +12,7 @@ from integration.translation.templates.vertex_ai import ( ) GEMINI_3_5_FLASH_THOUGHT_SIGNATURE: Final = "AY89a1+kriNs9YsPegJQxgjTWsMlxsqPU/UkHDyz4EK+3FRrE52ygqPIHTvkblgKKiO5Fn8R9D3iu//i+YnuuZ+e7W0yxUoFScas89b5rn3vZHsihnUA7VTe3b7beth1IftteoTGZ3F8OSLrSIP8HKn4xTIx35eeNEt5oiD3K5rTQrdklqO277HvTDhu/LuZhZ0pzaWciGQEc61epIBBpC2ygu8V4L8TekUVZnz6qzSzxipdiQStQ8dZou43334hMnxOR6JjJCGxqt/1K+C0vXX6YP6f+Bph5VhPSnvy2Z1uApi3TUxAcsDDyVewIr+ha2G6DPhN1Oi5uZZNhmbvH01Uq5OOquF1uEICq2m/35LhJIbDNM9N7oI1Z6L94P9xdugeFXM2S7wdJfc8sdg09vEzNeA++GK25zRuB4hU7XyaG4XnWjmu2kswi4sf+EyrO/K0P2tb/ZCJu7kJTtWTiD1wgSIl1eqS8ORKKsXDN2hIbFrtsNR50WIvVzd6" +CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE: Final = "CAQS0gYKEggSEAIYAjgBQgh0aGlua2luZxIMcbnLkNvj2D6NoBl7Ggw+Qgax2/68TAR1maEiMH6XEeKEqR1ewAiuj01fMNuD1KI7pCTXUsWWr6H7nSzP5BiRGXRrDCIxPYu2Xbvx7yrtBbbyId4ReZzx21wZ2+d0vkkpZM82PXY72wQRJOHLR4JNLRjiR2Jn+FrZ8mBo2mofFrlTSuioOVHc1hheiH4InRK57w2X9xaTXa+mwoWW5epPieHD26+NNJirF9vB832fvqDVOd+f8N0ZDUXUut8zl/Jxx+0BMOktMqyfehIOGumQmyv49E3JHVQGrx+Dd3KZjaQXrsRFf1K3AJ7p5XElR+cYis40W7EL3DE9EAtriZuDEt0fTct79fzZIZgHd6Dnua/rGRnjNjaaZKxGjnhfkVHwsZW4HqTNgwWQiD3lt2ncoLfsfd2IWwMEI5uWpN3GtvRkQvfzidmgGmx3dI8g0nBvLPuq0jzVIuuoKIKqzT7A5ae3Gs9eTlSDr2wOGfdyNf6Q2TLw0w7EegI4wYSExKhFkx0EytCGHJcbN0WClXXmHsRtrXnP7sxCMMY7B8TIpAoN7SXyNjJizOGgzoYL7PNUgqp5nVz/C9vTvudhSWbjLITdxFJ5hearL2rMhryfHBKvRGYC9RMy/EjjtINNwYpxdVX29mUE+8KT+41HymtJ1Ewjtoo8V/7XQnQk3cB7X2RFHcR89hAM/pSsB14ofMFoDk+jJitS+sCDZI0Xh7/ySU0o/OI4l+oD/yctZTHx/mIOiWRtnAOrveNT1ms0/2+4tcO0wi3pRrOB2s5ugYc1I3+GcRwILEIM3a4ABd5Dw2nW2fgUhfSTTGURm4Z91ciih8pKUS1ZPatFTOJ1ehKeUCXflkqY2PHnftiNB/8GzlomsQ2NoRw8UXD00oMmBvSvZpPN7khFLxWHtxLjleKkeKCzN7QZkOqvjAnUwy9TacT2ZvSXIxrvaep7dOyrlMEta/1KsYEcWuBxLr87Va4txXnmoV4RqGgaauNcR2wOGyRFJppxKAgYbmMnxams7cv/7IY2RVSgSQf2k4hFQYzRwrHngjapuZzgNS7UCMjvJJs4VdIypVOATsZdPzJToMgWaiEoRZ3zApQDgYhhGAE=" CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", @@ -790,6 +791,160 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/responses", + litellm_request={**RESPONSES_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request={ + **VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + "system": [ + { + "type": "text", + "text": ( + "You help maintain a synthetic Python project. Read the relevant files before suggesting a fix. " + "Explain which file you need next, then call a read-only tool. Never change files. Use the responses " + "client for this project review. Consider zero, negative and non-integer quantities before choosing " + "the next file to inspect." + ), + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8BhBrWmiX5ZtMxpM9zA", + "content": [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer values " + "before moving to the next file—I'll grab the first 60 lines since I don't know the full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read the validation function in `shop/validation.py`, starting around the match on line 18." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 255, + "output_tokens_details": {"thinking_tokens": 126}, + }, + }, + expected_litellm_response={ + **RESPONSES_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "output": [ + { + "type": "reasoning", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "annotations": [], + }, + ], + "phase": None, + "encrypted_content": json.dumps( + [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + ], + separators=(",", ":"), + ), + }, + { + "type": "message", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "I'll read the validation function in `shop/validation.py`, starting around the match on line 18.", + "annotations": [], + }, + ], + "phase": None, + }, + { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 60}', + "call_id": "toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "name": "read_file", + "type": "function_call", + "id": "fc_toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "namespace": None, + "status": "completed", + }, + ], + "usage": { + "input_tokens": 877, + "input_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "cached_tokens_details": None, + "image_tokens": None, + "text_tokens": 4, + "video_tokens": None, + "cache_write_tokens": 873, + }, + "output_tokens": 255, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 126, "text_tokens": 129}, + "total_tokens": 1132, + "cost": None, + }, + "provider_specific_fields": { + "citations": None, + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + ], + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/responses", diff --git a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py index 67c3cf4f7a9..bd9118482d0 100644 --- a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py +++ b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.responses.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 58549a028ea..388e4c42ea0 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -7,6 +7,7 @@ from dotenv import load_dotenv load_dotenv() import asyncio import hashlib +import threading import random import pytest @@ -14,6 +15,7 @@ from pydantic import TypeAdapter import litellm from litellm import aembedding, completion, embedding, aresponses, responses +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.caching.caching import Cache from litellm.responses.streaming_iterator import CachedResponsesAPIStreamingIterator @@ -65,6 +67,46 @@ def setup_cache(): return cache +@pytest.mark.parametrize( + ("call", "request_kwargs"), + [ + ( + completion, + {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "mock_response": "x"}, + ), + (embedding, {"model": "text-embedding-3-small", "input": ["hi"], "mock_response": [0.1, 0.2]}), + (litellm.text_completion, {"model": "gpt-3.5-turbo-instruct", "prompt": "hi", "mock_response": "x"}), + ], +) +def test_sync_cache_hits_reach_sync_custom_logger( + monkeypatch: pytest.MonkeyPatch, call: Callable[..., object], request_kwargs: dict[str, object] +) -> None: + from litellm.integrations.custom_logger import CustomLogger + + logged_cache_hits: Final[list[object]] = [] + all_logged: Final = threading.Event() + + class SyncRecorder(CustomLogger): + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + logged_cache_hits.append(kwargs.get("cache_hit")) + if len(logged_cache_hits) == 3: + all_logged.set() + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr(litellm, "callbacks", [SyncRecorder()]) + monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) + + for _ in range(3): + call(**request_kwargs) + + assert all_logged.wait(timeout=10) + assert logged_cache_hits == [None, True, True] + + @pytest.fixture def copilot_response_cache(monkeypatch: pytest.MonkeyPatch) -> Cache: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) @@ -262,6 +304,93 @@ async def test_cache_hit_reports_zero_response_cost(copilot_response_cache: Cach assert second._hidden_params["response_cost"] == 0 +@pytest.mark.asyncio +async def test_async_cache_hits_never_use_up_a_deployment_tpm_limit(copilot_response_cache: Cache) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "tpm-limited", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "mock_response": "cached answer", + "tpm": 100, + }, + } + ], + routing_strategy="usage-based-routing-v2", + num_retries=0, + ) + messages: Final = [{"role": "user", "content": "tpm limit and cache hits"}] + + async def call_and_settle() -> litellm.ModelResponse: + response: Final = await router.acompletion(model="tpm-limited", messages=messages, caching=True) + await asyncio.gather(*_PENDING_CACHE_WRITES) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + return response + + first: Final = await call_and_settle() + hits: Final = [await call_and_settle() for _ in range(8)] + + assert first.usage.total_tokens * len(hits) > 100 + assert all(hit._hidden_params["cache_hit"] is True for hit in hits) + + +def _logging_obj_left_with_a_cache_hit(messages: list[dict[str, str]]) -> LiteLLMLogging: + logging_obj: Final = LiteLLMLogging( + litellm_call_id=str(uuid.uuid4()), + call_type=CallTypes.acompletion.value, + model="gpt-4o-mini", + messages=messages, + function_id=str(uuid.uuid4()), + stream=False, + start_time=datetime.now(), + ) + logging_obj.caching_details = {"cache_hit": True, "cache_duration_ms": 1.0} + return logging_obj + + +@pytest.mark.asyncio +async def test_async_cache_miss_clears_the_hit_an_earlier_lookup_left_on_the_call( + copilot_response_cache: Cache, +) -> None: + messages: Final = [{"role": "user", "content": f"never cached {uuid.uuid4()}"}] + logging_obj: Final = _logging_obj_left_with_a_cache_hit(messages) + handler: Final = LLMCachingHandler( + original_function=litellm.acompletion, request_kwargs={}, start_time=datetime.now() + ) + + response: Final = await handler.async_get_cache( + model="gpt-4o-mini", + original_function=litellm.acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs={"model": "gpt-4o-mini", "messages": messages, "caching": True}, + ) + + assert response is not None and response.cached_result is None + assert logging_obj.caching_details is None + + +def test_sync_cache_miss_clears_the_hit_an_earlier_lookup_left_on_the_call(copilot_response_cache: Cache) -> None: + messages: Final = [{"role": "user", "content": f"never cached {uuid.uuid4()}"}] + logging_obj: Final = _logging_obj_left_with_a_cache_hit(messages) + handler: Final = LLMCachingHandler(original_function=completion, request_kwargs={}, start_time=datetime.now()) + + response: Final = handler.sync_get_cache( + model="gpt-4o-mini", + original_function=completion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.completion.value, + kwargs={"model": "gpt-4o-mini", "messages": messages, "caching": True}, + ) + + assert response.cached_result is None + assert logging_obj.caching_details is None + + chat_completion_response = litellm.ModelResponse( id=str(uuid.uuid4()), choices=[ @@ -1072,7 +1201,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): stream=True, start_time=datetime.now(), ) - logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + logging_obj.handle_sync_cache_hit_success_callbacks = MagicMock() kwargs = { "model": original_model, @@ -1096,7 +1225,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): assert isinstance( cached_response.cached_result, CachedResponsesAPIStreamingIterator ) - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + logging_obj.handle_sync_cache_hit_success_callbacks.assert_not_called() def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): @@ -1116,7 +1245,7 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): stream=True, start_time=datetime.now(), ) - logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + logging_obj.handle_sync_cache_hit_success_callbacks = MagicMock() kwargs = { "model": original_model, @@ -1137,7 +1266,7 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): ) assert cached_response.cached_result is not None - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + logging_obj.handle_sync_cache_hit_success_callbacks.assert_not_called() def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request(): @@ -3069,7 +3198,7 @@ async def test_a_cache_hit_inside_an_internal_sub_call_is_free_and_fires_no_call async_logging.async_success_handler = AsyncMock() async_logging.handle_sync_success_callbacks_for_async_calls = MagicMock() sync_logging: Final = _build_logging_obj(CallTypes.completion.value, stream=False) - sync_logging.handle_sync_success_callbacks_for_async_calls = MagicMock() + sync_logging.handle_sync_cache_hit_success_callbacks = MagicMock() with internal_sub_call() if internal else nullcontext(): async_hit: Final = await LLMCachingHandler( @@ -3100,4 +3229,4 @@ async def test_a_cache_hit_inside_an_internal_sub_call_is_free_and_fires_no_call assert sync_hit.cached_result._hidden_params.get("response_cost") == 0.0 assert async_logging.async_success_handler.called is not internal assert async_logging.handle_sync_success_callbacks_for_async_calls.called is not internal - assert sync_logging.handle_sync_success_callbacks_for_async_calls.called is not internal + assert sync_logging.handle_sync_cache_hit_success_callbacks.called is not internal diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 7efcc76073f..8331fb1937b 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,11 +2,11 @@ import asyncio import base64 import importlib import os -from collections.abc import Coroutine, Iterator +from collections.abc import Callable, Coroutine, Iterator from dataclasses import dataclass, field from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias import boto3 import httpx @@ -308,6 +308,22 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +NativeStructuredOutputKey: TypeAlias = Literal["legacy", "bedrock"] | None + + +@pytest.fixture +def set_native_structured_output( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> Callable[[str, NativeStructuredOutputKey], None]: + def set_flags(model_cost_key: str, native_key: NativeStructuredOutputKey) -> None: + entry: Final = litellm.model_cost[model_cost_key] + monkeypatch.setitem(entry, "supports_native_structured_output", native_key == "legacy") + monkeypatch.setitem(entry, "supports_bedrock_native_structured_output", native_key == "bedrock") + litellm.get_model_info.cache_clear() + + return set_flags + + @pytest.fixture def local_beta_headers_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 8b1a536d8e3..20af774278f 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -24,6 +24,7 @@ from litellm.constants import AWS_SIGNING_MAX_THREADS from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.decisions import ( @@ -120,6 +121,7 @@ _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( "pplx-decider-v1-27b", ), ("typesafe", "typesafe/jev-1.13", "https://api.typesafe.ai/v1/systemone", "jev-1.13"), + ("nace", "nace/drex-v2.0", "https://console.nace.ai/v1/systemone", "drex-v2.0"), ( "openrouter", "openrouter/typesafe/jev-1.13", @@ -2681,6 +2683,269 @@ async def test_azure_ai_decision_requires_api_base_before_http( assert len(respx_mock.calls) == 0 +def _clear_azure_ai_auth_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) + monkeypatch.delenv("AZURE_TENANT_ID", raising=False) + monkeypatch.delenv("AZURE_CLIENT_ID", raising=False) + monkeypatch.delenv("AZURE_CLIENT_SECRET", raising=False) + monkeypatch.delenv("AZURE_USERNAME", raising=False) + monkeypatch.delenv("AZURE_PASSWORD", raising=False) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) + + +def test_azure_ai_systemone_uses_entra_token_without_api_key( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + litellm.systemone( + model="azure_ai/decision-1", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + api_base="https://res.services.ai.azure.com", + azure_ad_token="entra-access-token", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer entra-access-token" + + +@pytest.mark.asyncio +async def test_azure_ai_decisions_use_entra_token_without_api_key( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + await litellm.adecisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + azure_ad_token="entra-access-token", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer entra-access-token" + + +def test_azure_ai_systemone_environment_api_key_wins_over_deployment_entra( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + monkeypatch.setenv("AZURE_AI_API_KEY", "env-key") + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + litellm.systemone( + model="azure_ai/decision-1", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + api_base="https://res.services.ai.azure.com", + azure_ad_token="entra-access-token", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer env-key" + + +@pytest.mark.asyncio +async def test_azure_ai_decisions_environment_api_key_wins_over_deployment_entra( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + monkeypatch.setenv("AZURE_AI_API_KEY", "env-key") + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + await litellm.adecisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + azure_ad_token="entra-access-token", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer env-key" + + +def test_azure_ai_decisions_uses_environment_api_key_without_entra_credentials( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + monkeypatch.setenv("AZURE_AI_API_KEY", "env-key") + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + litellm.decisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer env-key" + + +def test_azure_ai_decisions_explicit_api_key_precedes_entra_and_environment( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + monkeypatch.setenv("AZURE_AI_API_KEY", "env-key") + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + litellm.decisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + api_key="call-key", + azure_ad_token="entra-access-token", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer call-key" + + +def test_azure_ai_decisions_does_not_validate_noncredential_params_as_auth( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + litellm.decisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + azure_ad_token="entra-access-token", + max_retries="three", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer entra-access-token" + + +def test_azure_ai_decisions_malformed_credential_is_a_sanitized_bad_request( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + with pytest.raises(litellm.BadRequestError) as caught: + litellm.decisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + tenant_id="tenant-id", + client_id="client-id", + client_secret=123, + ) + + assert caught.value.status_code == 400 + assert "`client_secret`" in str(caught.value) + assert "123" not in str(caught.value) + assert not route.called + + +def test_azure_ai_decisions_prefer_api_key_when_both_are_set( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + litellm.decisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + api_key="foundry-key", + azure_ad_token="entra-access-token", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer foundry-key" + + +def test_azure_ai_decisions_without_credentials_raise_authentication_error( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + with pytest.raises(ValueError, match="Missing Azure AI credentials") as helper_error: + get_azure_ai_auth_headers(None, {}) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + with pytest.raises(litellm.AuthenticationError) as caught: + litellm.decisions( + model="azure_ai/decision-1", + input="review", + questions=_PREDICATE_QUESTIONS, + api_base="https://res.services.ai.azure.com", + ) + + assert caught.value.status_code == 401 + assert str(helper_error.value) in str(caught.value) + assert not route.called + + +@pytest.mark.asyncio +async def test_azure_ai_systemone_without_credentials_raises_authentication_error( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + _clear_azure_ai_auth_environment(monkeypatch) + with pytest.raises(ValueError, match="Missing Azure AI credentials") as helper_error: + get_azure_ai_auth_headers(None, {}) + route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond( + json={**_RESPONSE, "model": "microsoft-decision-1"} + ) + + with pytest.raises(litellm.AuthenticationError) as caught: + await litellm.asystemone( + model="azure_ai/decision-1", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + api_base="https://res.services.ai.azure.com", + ) + + assert caught.value.status_code == 401 + assert str(helper_error.value) in str(caught.value) + assert not route.called + + _CLEF_OMNI_UPSTREAM: Final = "https://example.com/ai/run/@cf/cloudflare/clef-omni" _CLEF_OMNI_RESPONSE: Final[Mapping[str, object]] = { "success": True, @@ -3345,3 +3610,137 @@ def test_dropping_every_question_answers_400_before_http(respx_mock: respx.MockR assert caught.value.status_code == 400 assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("openai_format", (False, True)) +@pytest.mark.parametrize( + ("provider", "model", "url", "key_env"), + ( + ("nace", "drex-test-alias", "https://console.nace.ai/v1/systemone", "NACE_API_KEY"), + ("openrouter", "nace/drex-test-alias", "https://openrouter.ai/api/alpha/decisions", "OPENROUTER_API_KEY"), + ( + "fireworks_ai", + "accounts/test/models/drex-test-alias", + "https://api.fireworks.ai/inference/v1/systemone", + "FIREWORKS_API_KEY", + ), + ), +) +async def test_decision_hosts_preserve_drex_model_routing_credentials_and_returned_model_pricing( + provider: str, + model: str, + url: str, + key_env: str, + openai_format: bool, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv(key_env, "host-test-key") + monkeypatch.delenv("NACE_API_BASE", raising=False) + monkeypatch.delenv("NACE_BASE_URL", raising=False) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + monkeypatch.delenv("FIREWORKS_API_BASE", raising=False) + monkeypatch.setitem( + litellm.model_cost, + f"{provider}/{model}", + { + "litellm_provider": provider, + "mode": "evaluation", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0, + }, + ) + monkeypatch.setitem( + litellm.model_cost, + f"{provider}/drex-test-resolved", + { + "litellm_provider": provider, + "mode": "evaluation", + "input_cost_per_token": 0.002, + "output_cost_per_token": 0, + }, + ) + payload: Final = {**_RESPONSE, "model": "drex-test-resolved", "request_id": "request-test"} + upstream: Final = respx_mock.post(url).respond(json=payload) + if openai_format: + response: Final = await litellm.adecisions( + model=f"{provider}/{model}", + input="damaged order", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + ) + assert isinstance(response, OpenAIDecisionResponse) + assert response.answers[0].model_dump() == {"type": "predicate", "name": "is_defect", "probability": 0.9} + else: + systemone_response: Final = await litellm.asystemone( + model=f"{provider}/{model}", + state="damaged order", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + assert systemone_response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert litellm.completion_cost(completion_response=systemone_response, call_type="asystemone") == pytest.approx( + _INPUT_TOKENS * 0.002 + ) + assert upstream.call_count == 1 + assert upstream.calls[0].request.headers["authorization"] == "Bearer host-test-key" + assert json.loads(upstream.calls[0].request.content) == { + "model": model, + "state": "damaged order", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + + +@pytest.mark.parametrize("base_env", ("NACE_API_BASE", "NACE_BASE_URL")) +def test_nace_sync_call_uses_environment_base_and_explicit_credentials( + base_env: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv(base_env, "https://nace.example/v1/") + monkeypatch.setenv("NACE_API_KEY", "environment-key") + upstream: Final = respx_mock.post("https://nace.example/v1/systemone").respond(json=_RESPONSE) + response: Final = litellm.systemone( + model="nace/drex-v2.0", + state={"message": "damaged order"}, + api_key="explicit-key", + questions={"is_defect": {"type": "noul", "criteria": {"true": "damaged", "false": "intact"}}}, + ) + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert upstream.calls[0].request.headers["authorization"] == "Bearer explicit-key" + assert json.loads(upstream.calls[0].request.content) == { + "model": "drex-v2.0", + "state": {"message": "damaged order"}, + "questions": {"is_defect": {"type": "noul", "criteria": {"true": "damaged", "false": "intact"}}}, + } + + +@pytest.mark.parametrize( + ("status", "exception"), + ( + (401, litellm.AuthenticationError), + (402, litellm.PaymentRequiredError), + (422, litellm.BadRequestError), + (429, litellm.RateLimitError), + (529, litellm.APIError), + ), +) +def test_nace_upstream_failures_preserve_status_and_message( + status: int, + exception: type[Exception], + respx_mock: respx.MockRouter, +) -> None: + upstream: Final = respx_mock.post("https://console.nace.ai/v1/systemone").respond( + status, + json={"error": {"message": "decision rejected"}, "request_id": "request-test"}, + headers={"retry-after": "1"}, + ) + with pytest.raises(exception, match="decision rejected") as error: + litellm.systemone( + model="nace/drex-v2.0", + state="text", + api_key="test-key", + questions={"a": {"type": "noul", "instructions": "Is it urgent?"}}, + num_retries=0, + ) + assert getattr(error.value, "status_code") == status + assert upstream.call_count == 1 diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 2106436edb8..4b053fb5a97 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -20,6 +20,7 @@ import pytest import yaml from pydantic import TypeAdapter +from litellm.proxy.types_utils.utils import get_instance_fn from tests.integration._support.client import Gateway, eventually from tests.integration._support.process import ( OwnedProxy, @@ -36,6 +37,32 @@ from tests.integration._support.process import ( TESTS_DIR: Final = Path(__file__).resolve().parents[2] +@pytest.mark.parametrize("config_contents", ("[]\n", "general_settings: []\n", "general_settings: {}\n")) +def test_worker_configuration_preserves_sibling_auth_modules( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + config_contents: str, +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203") + monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202") + source: Final = tmp_path / "source.yaml" + source.write_text(config_contents) + auth_state: Final = tmp_path / "sibling_auth.json" + auth_state.write_text("authenticated") + (tmp_path / "sibling_auth.py").write_text( + "from pathlib import Path\n" + "def authenticate():\n" + " return Path(__file__).with_suffix('.json').read_text()\n" + ) + + configured: Final = worker_configuration(source, tmp_path / "worker") + + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "authenticated" + auth_state.write_text("updated") + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "updated" + + def test_worker_configuration_routes_stock_edges_to_its_own_servers_and_preserves_other_targets( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: @@ -93,6 +120,42 @@ def test_worker_configuration_budgets_implicit_pools_without_changing_connection assert yaml.safe_load(source.read_text()) == original +def test_worker_configuration_keeps_modules_beside_a_test_config_importable( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + source: Final = tmp_path / "proxy.yaml" + source.write_text(yaml.safe_dump({"general_settings": {"custom_auth": "sibling_auth.authenticate"}})) + (tmp_path / "sibling_auth.py").write_text("def authenticate() -> str:\n return 'sibling'\n") + configured: Final = worker_configuration(source, tmp_path / "gateway") + assert configured != source + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "sibling" + + +def test_worker_configuration_writes_the_repository_config_copy_into_the_worker_directory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203") + monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202") + monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(tmp_path / "product")) + configured: Final = worker_configuration(TESTS_DIR / "integration/proxy_config.yaml", tmp_path / "worker") + assert configured.parent == tmp_path / "worker" + + +def test_worker_configuration_writes_the_proxy_checkout_config_copy_into_the_worker_directory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + checkout: Final = tmp_path / "checkout" + monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(checkout)) + source: Final = checkout / "tests/integration/proxy_config.yaml" + source.parent.mkdir(parents=True) + source.write_text("model_list: []\n") + configured: Final = worker_configuration(source, tmp_path / "worker") + assert configured.parent == tmp_path / "worker" + + def test_serial_proxy_keeps_the_requested_configuration(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: monkeypatch.delenv("INTEGRATION_WORKER_ISOLATION", raising=False) source: Final = tmp_path / "source.yaml" diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index bc230a8bcb2..3f4389765d3 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -2,6 +2,7 @@ baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name builders, and the registry validator's failure paths. Needs the OTel SDK.""" +from typing import Final import contextlib import json import threading @@ -358,6 +359,18 @@ def test_genai_mapper_cost_breakdown_absent(): assert not any(k.startswith(LiteLLM.COST_PREFIX) and k != f"{LiteLLM.COST_PREFIX}total" for k in attrs) +def test_genai_mapper_cache_hit_attrs(): + from litellm.integrations.otel.model.semconv import LiteLLM + + attrs: Final = GenAIMapper().map(replace(_full_llm_call(), cache_hit=True, saved_cache_cost=0.002)) + assert attrs[LiteLLM.CACHE_HIT] is True + assert attrs[LiteLLM.SAVED_CACHE_COST] == 0.002 + + uncached: Final = GenAIMapper().map(_full_llm_call()) + assert LiteLLM.CACHE_HIT not in uncached + assert LiteLLM.SAVED_CACHE_COST not in uncached + + def test_llm_cost_from_breakdown_maps_costbreakdown_keys(): cost = LLMCost.from_breakdown( { diff --git a/tests/unit/integrations/otel/test_otel_v2_emitter.py b/tests/unit/integrations/otel/test_otel_v2_emitter.py index 949eb97bccd..0905e9e7c30 100644 --- a/tests/unit/integrations/otel/test_otel_v2_emitter.py +++ b/tests/unit/integrations/otel/test_otel_v2_emitter.py @@ -2,6 +2,7 @@ legacy dual-emit, hierarchy, error status, and idempotency. Needs the OTel SDK.""" import json +from typing import Final import pytest @@ -642,25 +643,71 @@ def test_preset_message_key_the_fit_sheds_still_never_overflows_the_span(): assert len(finished.attributes) <= SpanLimits().max_span_attributes -def test_attribute_budget_counts_only_keys_the_fit_does_not_overwrite(): +def test_preset_keys_the_mappers_overwrite_shed_no_message_attribute(monkeypatch): """Pre-set attributes the mapped set overwrites consume no slot against the span limit. A boundary-opened LLM-call span already carries a few stamped attributes, - most of which the mappers re-emit under the same keys. Charging those - against the budget reserves slots the fit can never spend and sheds - indexed message attributes for nothing. + most of which the mappers re-emit under the same keys, so a span sized to + exactly its mapped attributes keeps every indexed message. """ + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "1000") + full: Final = _conversation_span(["genai", "openinference"], _conversation_payload(6)).attributes + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", str(len(full))) + cfg: Final = OpenTelemetryV2Config( + exporter="in_memory", + legacy_compat=False, + mapper_names=["genai", "openinference"], + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine: Final = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + span: Final = engine.start_span(SpanRole.LLM_CALL, "chat gpt-4o") + span.set_attribute("gen_ai.request.model", "gpt-4o") + span.set_attribute("llm.input_messages.1.message.role", "stale-role") + engine.finish_span( + SpanRole.LLM_CALL, + span, + LLMCallSpanData.from_standard_logging_payload(_conversation_payload(6), capture_content=True), + ) + (finished,) = exporter.get_finished_spans() + assert finished.dropped_attributes == 0 + assert dict(finished.attributes) == dict(full) + + +def test_a_preset_message_key_the_fit_sheds_never_overflows_an_exact_limit(monkeypatch): + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "1000") + full: Final = _conversation_span(["genai", "openinference"], _conversation_payload(6)).attributes + limit: Final = len(full) - 2 + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", str(limit)) + cfg: Final = OpenTelemetryV2Config( + exporter="in_memory", + legacy_compat=False, + mapper_names=["genai", "openinference"], + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine: Final = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + span: Final = engine.start_span(SpanRole.LLM_CALL, "chat gpt-4o") + span.set_attribute("llm.input_messages.1.message.role", "stale-role") + engine.finish_span( + SpanRole.LLM_CALL, + span, + LLMCallSpanData.from_standard_logging_payload(_conversation_payload(6), capture_content=True), + ) + (finished,) = exporter.get_finished_spans() + assert finished.dropped_attributes == 0 + assert len(finished.attributes) <= limit + assert finished.attributes.get("llm.input_messages.1.message.role") != "stale-role" + + +def test_attribute_budget_charges_every_key_the_span_already_carries(): cfg = OpenTelemetryV2Config(exporter="in_memory") provider, _exporter = providers.in_memory_provider(cfg) span = providers.get_tracer(provider, "litellm-test").start_span("s") span.set_attribute("gen_ai.request.model", "gpt-4o") span.set_attribute("litellm.metadata.user_api_key_alias", "edge-key") try: - assert attribute_budget(span, 0) == SpanLimits().max_span_attributes - 2 - assert ( - attribute_budget(span, 0, frozenset({"gen_ai.request.model"})) - == SpanLimits().max_span_attributes - 1 - ) + assert attribute_budget(span, 1) == SpanLimits().max_span_attributes - 3 finally: span.end() diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 6a9a963948e..10e12981ec1 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -494,6 +494,49 @@ def test_no_span_when_pre_call_never_ran(): assert exporter.get_finished_spans() == () # no phantom LLM span +def test_cache_hit_without_pre_call_emits_zero_cost_span(): + """A response served from the litellm cache never hands off to a provider, + so ``pre_call`` never runs and ``api_call_start_time`` is never stamped, + but the completed call's payload still warrants a span, carrying the cache + attributes rather than being dropped like a gate rejection.""" + logger, exporter = _logger() + payload: Final = _payload( + cache_hit=True, + response_cost=0.0, + saved_cache_cost=0.002, + cost_breakdown={"input_cost": 0.0, "output_cost": 0.0, "total_cost": 0.0, "tool_usage_cost": 0.0}, + ) + asyncio.run( + logger.async_log_success_event(_kwargs(payload=payload), None, None, None) + ) + (span,) = exporter.get_finished_spans() + assert span.attributes[LiteLLM.CACHE_HIT] is True + assert span.attributes[LiteLLM.SAVED_CACHE_COST] == 0.002 + assert { + name: span.attributes[f"{LiteLLM.COST_PREFIX}{name}"] for name in ("total", "input", "output", "tool_usage") + } == {"total": 0.0, "input": 0.0, "output": 0.0, "tool_usage": 0.0} + assert "server.address" not in span.attributes + + +def test_cache_miss_span_reports_cache_hit_false(): + logger, exporter = _logger() + _emit_llm(logger, _kwargs(payload=_payload(cache_hit=False, saved_cache_cost=0.0))) + (span,) = exporter.get_finished_spans() + assert span.attributes[LiteLLM.CACHE_HIT] is False + assert LiteLLM.SAVED_CACHE_COST not in span.attributes + + +def test_cache_miss_without_pre_call_still_emits_no_span(): + """The cache-hit relaxation is not a general loosening: a success callback + with no ``pre_call`` carrier and no cache hit still emits nothing.""" + logger, exporter = _logger() + payload: Final = _payload(cache_hit=False, response_cost=0.002) + asyncio.run( + logger.async_log_success_event(_kwargs(payload=payload), None, None, None) + ) + assert exporter.get_finished_spans() == () + + def test_real_llm_failure_still_emitted(): """A genuine LLM failure: ``pre_call`` ran (the call was attempted), so the CLIENT span is opened at the boundary and closed ERROR.""" diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 9ba09085ce8..b40c1a96500 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -478,6 +478,75 @@ class TestOpenTelemetryCostBreakdown(unittest.TestCase): assert ("gen_ai.cost.tool_usage_cost", 0.0001) not in call_args_list assert ("gen_ai.cost.original_cost", 0.004) not in call_args_list + def test_cache_hit_span_carries_cache_indicator_and_saved_cost(self): + """ + A cache-hit span reports litellm.cache_hit and litellm.saved_cache_cost so cost + attribution can tell a cache-served response apart from a billed call. + """ + otel: Final = OpenTelemetry() + mock_span: Final = MagicMock() + + kwargs: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "cache_hit": True, + "saved_cache_cost": 0.003, + "cost_breakdown": {"input_cost": 0.0, "output_cost": 0.0, "total_cost": 0.0, "tool_usage_cost": 0.0}, + }, + } + + response_obj: Final = { + "id": "test-response-id", + "model": "gpt-4", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj=response_obj) + + mock_span.set_attribute.assert_any_call("litellm.cache_hit", True) + mock_span.set_attribute.assert_any_call("litellm.saved_cache_cost", 0.003) + mock_span.set_attribute.assert_any_call("gen_ai.cost.total_cost", 0.0) + mock_span.set_attribute.assert_any_call("gen_ai.cost.input_cost", 0.0) + mock_span.set_attribute.assert_any_call("gen_ai.cost.output_cost", 0.0) + + def test_non_cache_hit_span_reports_cache_hit_false(self): + otel: Final = OpenTelemetry() + mock_span: Final = MagicMock() + + kwargs: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "cache_hit": False, + }, + } + + response_obj: Final = { + "id": "test-response-id", + "model": "gpt-4", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj=response_obj) + + call_args_list: Final = [call[0] for call in mock_span.set_attribute.call_args_list] + assert ("litellm.cache_hit", False) in call_args_list + assert all(key != "litellm.saved_cache_cost" for key, _value in call_args_list) + assert all(not key.startswith("gen_ai.cost.") for key, _value in call_args_list) + class TestOpenTelemetryProviderInitialization(unittest.TestCase): """Test suite for verifying provider initialization respects existing providers""" diff --git a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py index 45b378d10fe..5e5a32575b9 100644 --- a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py +++ b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py @@ -20,10 +20,14 @@ Tests cover: - llm_router unavailable / model_group missing / router raises → silent no-op. """ +from collections.abc import Iterator +from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from prometheus_client.samples import Sample from litellm.integrations.prometheus import PrometheusLogger @@ -293,3 +297,89 @@ class TestRouterFallbackDefensivePaths: ) prometheus_logger.litellm_remaining_tokens_metric.labels.assert_not_called() + + +def _registry_samples() -> Iterator[Sample]: + for metric in REGISTRY.collect(): + yield from metric.samples + + +def _sample_values(sample_name: str) -> tuple[float, ...]: + return tuple(sample.value for sample in _registry_samples() if sample.name == sample_name) + + +@pytest.mark.parametrize("remaining_tokens", ["856, abc", 10**400], ids=["garbled", "too-large-for-a-float"]) +def test_unparseable_provider_ratelimit_header_skips_only_its_own_gauge( + prometheus_logger: PrometheusLogger, remaining_tokens: str | int +): + payload: Final = _build_payload( + model_group="groq-group", + custom_llm_provider="groq", + additional_headers={ + "x_ratelimit_remaining_requests": 0, + "x_ratelimit_remaining_tokens": remaining_tokens, + }, + ) + request_kwargs: Final = { + "litellm_params": {"custom_llm_provider": "groq", "metadata": {}}, + "model": "llama-3.1-8b-instant", + "standard_logging_object": payload, + } + + prometheus_logger.set_llm_deployment_success_metrics( + request_kwargs=request_kwargs, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1, 0, 0, 1), + enum_values=_enum_values(model_group="groq-group"), + output_tokens=2.0, + ) + + assert _sample_values("litellm_remaining_requests_metric") == (0.0,) + assert _sample_values("litellm_remaining_tokens_metric") == () + assert _sample_values("litellm_deployment_success_responses_total") == (1.0,) + + +@pytest.mark.parametrize( + ("additional_headers", "remaining_requests", "remaining_tokens"), + [ + ({"x_ratelimit_remaining_requests": "12.5", "x_ratelimit_remaining_tokens": "856.5"}, 12.5, 856.5), + ( + {1: "non-string key", "x_ratelimit_remaining_requests": "12.5", "x_ratelimit_remaining_tokens": "856.5"}, + 12.5, + 856.5, + ), + ( + {"x_ratelimit_remaining_requests": "inf", "x_ratelimit_remaining_tokens": "\u0661\u0662.\u0665"}, + float("inf"), + 12.5, + ), + ], + ids=["decimal", "decimal-with-a-non-string-key", "infinity-and-arabic-indic-digits"], +) +def test_provider_ratelimit_headers_set_both_gauges_to_the_value_float_reads( + prometheus_logger: PrometheusLogger, + additional_headers: dict[object, object], + remaining_requests: float, + remaining_tokens: float, +): + payload: Final = _build_payload( + model_group="groq-group", + custom_llm_provider="groq", + additional_headers=additional_headers, + ) + + prometheus_logger.set_llm_deployment_success_metrics( + request_kwargs={ + "litellm_params": {"custom_llm_provider": "groq", "metadata": {}}, + "model": "llama-3.1-8b-instant", + "standard_logging_object": payload, + }, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1, 0, 0, 1), + enum_values=_enum_values(model_group="groq-group"), + output_tokens=2.0, + ) + + assert _sample_values("litellm_remaining_requests_metric") == (remaining_requests,) + assert _sample_values("litellm_remaining_tokens_metric") == (remaining_tokens,) + assert _sample_values("litellm_deployment_success_responses_total") == (1.0,) diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 6b017f9ff79..348d0105a6c 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -155,6 +155,89 @@ async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_ assert rebuilt.usage.completion_tokens == 7 +def _prompt_cached_payload(content: str, cache_read: int, cache_write: int) -> dict: + cached = _async_payload(content=content) + cached["usage"] = litellm.Usage( + prompt_tokens=42145, + completion_tokens=6, + total_tokens=42151, + cache_read_input_tokens=cache_read, + cache_creation_input_tokens=cache_write, + ).model_dump() + return cached + + +@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"]) +@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_the_cached_prompt_cache_tokens(replay, cache_read, cache_write, content): + cached = _prompt_cached_payload(content, cache_read, cache_write) + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None) + + assert usage.prompt_tokens == 42145 + assert usage.prompt_tokens_details.cached_tokens == cache_read + assert usage.prompt_tokens_details.cache_write_tokens == cache_write + assert usage.cache_read_input_tokens == cache_read + assert usage.cache_creation_input_tokens == cache_write + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_cached_tokens_reported_only_in_the_prompt_tokens_details(replay): + cached = _async_payload() + cached["usage"]["prompt_tokens_details"] = {"cached_tokens": 4, "audio_tokens": 1} + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.usage.prompt_tokens_details.cached_tokens == 4 + assert rebuilt.usage.prompt_tokens_details.audio_tokens == 1 + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_integral_counts_cached_as_floats_or_numeric_strings(replay): + cached = _async_payload() + cached["usage"] = { + "prompt_tokens": 100.0, + "completion_tokens": "6", + "total_tokens": 106.0, + "cache_read_input_tokens": 90.0, + } + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None) + + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (100, 6, 106) + assert usage.cache_read_input_tokens == 90 + + +@pytest.mark.parametrize( + "field, value", + [ + ("prompt_tokens_details", {"cached_tokens": "abc"}), + ("prompt_tokens_details", [1, 2]), + ("cache_read_input_tokens", "abc"), + ("cache_creation_input_tokens", 1.5), + ], + ids=["details_text", "details_list", "read_text", "write_fraction"], +) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_skips_cached_prompt_cache_usage_that_does_not_validate(replay, field, value): + cached = _async_payload() + cached["usage"][field] = value + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.content == "Hello world! How are you?" + assert rebuilt.usage.prompt_tokens == 5 + + @pytest.mark.asyncio async def test_async_finish_reason_only_on_last_chunk(): chunks = await _collect_async(_async_payload()) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1fda656060a..c08eeea4bb9 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -4,6 +4,7 @@ import functools import json import os import sys +from collections.abc import Mapping from typing import Final, cast import pytest @@ -29,6 +30,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( parse_tool_call_arguments, responses_reasoning_items_from_thinking_blocks, salvage_concatenated_tool_arguments, + schema_grammar_branch_count, + schema_has_open_object, + schema_has_unresolvable_ref, + schema_property_count, split_concatenated_json_objects, strip_encrypted_reasoning_from_messages, strip_litellm_response_only_fields, @@ -2362,3 +2367,75 @@ def test_strip_litellm_response_only_fields_drops_nested_tool_call_provider_fiel assert "provider_specific_fields" not in function assert function["name"] == "get_weather" assert function["arguments"] == '{"city": "Paris"}' + + +@pytest.mark.parametrize( + ("schema", "is_open"), + [ + ({"type": "object", "properties": {"a": {"type": "string"}}}, False), + ({"type": "object", "properties": {}}, False), + ({"type": "object", "additionalProperties": False}, False), + ({"type": "string"}, False), + ({"type": "object"}, True), + ({"type": ["object", "null"]}, True), + ({"type": "object", "properties": {"a": {"type": "string"}}, "additionalProperties": True}, True), + ({"type": "object", "additionalProperties": {"type": "integer"}}, True), + ({"type": "object", "properties": {}, "patternProperties": {"^x": {"type": "string"}}}, True), + ({"type": "object", "properties": {"inner": {"anyOf": [{"type": "object"}, {"type": "null"}]}}}, True), + ({"type": "array", "items": {"type": "object", "additionalProperties": {"type": "string"}}}, True), + ], +) +def test_schema_has_open_object_flags_any_node_that_accepts_unlisted_keys(schema, is_open): + assert schema_has_open_object(schema) is is_open + + +def _object_with(properties: Mapping[str, object], **defs: Mapping[str, object]) -> Mapping[str, object]: + return {"type": "object", "properties": properties, **defs} + + +@pytest.mark.parametrize( + ("schema", "unresolvable"), + [ + (_object_with({"a": {"type": "string"}}), False), + (_object_with({"a": {"$ref": "#/$defs/A"}, "b": {"$ref": "#/$defs/A"}}, **{"$defs": {"A": {"type": "string"}}}), False), + (_object_with({"a": {"$ref": "#/definitions/A"}}, definitions={"A": {"$ref": "#/definitions/B"}, "B": {"type": "integer"}}), False), + (_object_with({"a": {"$ref": "/$defs/A"}}, **{"$defs": {"A": {"items": {"$ref": "/$defs/B"}}, "B": {"type": "string"}}}), False), + (_object_with({"a": {"$ref": "#/$defs/Missing"}}, **{"$defs": {"A": {"type": "string"}}}), True), + (_object_with({"a": {"$ref": "other.json#/A"}}), True), + (_object_with({"parent": {"$ref": "#"}}), True), + (_object_with({"a": {"$ref": "#/$defs/A"}}, **{"$defs": {"A": {"type": "array", "items": {"$ref": "#/$defs/A"}}}}), True), + (_object_with({"a": {"$ref": "/$defs/A"}}, **{"$defs": {"A": {"type": "array", "items": {"$ref": "/$defs/A"}}}}), True), + ( + _object_with( + {"a": {"$ref": "#/$defs/A"}}, + **{"$defs": {"A": {"properties": {"b": {"$ref": "#/$defs/B"}}}, "B": {"properties": {"a": {"$ref": "#/$defs/A"}}}}}, + ), + True, + ), + ], + ids=["no-refs", "shared-def", "chained-definitions", "root-path-defs", "missing-def", "external", "root", "self-cycle", "root-path-self-cycle", "mutual-cycle"], +) +def test_schema_has_unresolvable_ref_flags_refs_that_cannot_be_inlined(schema, unresolvable): + assert schema_has_unresolvable_ref(schema) is unresolvable + + +@pytest.mark.parametrize( + ("schema", "properties", "branches"), + [ + (_object_with({"a": {"type": "string"}, "b": {"type": "string"}, "c": {"type": "string"}}, required=["a"]), 3, 8), + ( + _object_with( + {"x": {"$ref": "#/$defs/Pair"}, "y": {"$ref": "#/$defs/Pair"}}, + required=["x", "y"], + **{"$defs": {"Pair": _object_with({"l": {"type": "string"}, "r": {"type": "string"}}, required=["l"])}}, + ), + 4, + 7, + ), + (_object_with({"inner": _object_with({"z": {"type": "string"}})}, required=["inner"]), 2, 4), + ], + ids=["flat", "shared-def-counted-once", "nested"], +) +def test_schema_size_counts_every_written_object_once(schema, properties, branches): + assert schema_property_count(schema) == properties + assert schema_grammar_branch_count(schema) == branches diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index bf977325a85..ec6af7f27c3 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1892,6 +1892,25 @@ def test_bedrock_tools_unpack_defs(): _bedrock_tools_pt(tools=tools) +@pytest.mark.parametrize( + "parameters", + [ + {"type": ["object", "null"]}, + {"type": "object", "properties": {"tier": {"type": "integer"}}, "$defs": [{"type": "string"}]}, + ], + ids=["root_type_list", "defs_list"], +) +def test_bedrock_tools_pt_sends_malformed_parameters_as_an_object_schema(parameters: dict[str, object]) -> None: + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt + + tools: Final = [{"type": "function", "function": {"name": "json_tool_call", "parameters": parameters}}] + + input_schema: Final = _bedrock_tools_pt(tools=tools)[0]["toolSpec"]["inputSchema"]["json"] + + assert input_schema["type"] == "object" + assert "$defs" not in input_schema + + def test_bedrock_tools_pt_strict_parameter(): """Regression for strict tools on the Bedrock Converse path. diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index e891706df13..3cbc3591009 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -25,13 +25,13 @@ from typing import Any, Dict, Final, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch import pytest - +from openai._legacy_response import LegacyAPIResponse import litellm -from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.chat_completion_agentic_loop import ( maybe_run_chat_completion_agentic_loop, ) @@ -45,9 +45,9 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, ) from litellm.types.utils import ( + ChatCompletionMessageToolCall, Choices, Function, - ChatCompletionMessageToolCall, Message, ModelResponse, ) @@ -142,7 +142,7 @@ def _raw_response_for(model_response: ModelResponse) -> MagicMock: (an object exposing `.headers` and `.parse()` -> something with model_dump).""" parsed = MagicMock() parsed.model_dump.return_value = model_response.model_dump() - raw = MagicMock() + raw = MagicMock(spec=LegacyAPIResponse) raw.headers = {} raw.parse.return_value = parsed return raw @@ -171,7 +171,7 @@ async def test_internal_control_fields_never_leak_into_provider_body(restore_cal ] ) mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create = create + mock_client.post = create response = await litellm.acompletion( model="openai/gpt-4o-mini", diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 9a7f96000e6..c3a2f032651 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -1,6 +1,8 @@ """Tests for litellm_core_utils.core_helpers module.""" import logging +import threading +from collections.abc import Iterator, Mapping from typing import Final import httpx @@ -13,6 +15,7 @@ from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, + get_litellm_metadata_from_kwargs, get_parent_otel_span_from_kwargs, get_or_create_metadata_bucket, get_provider_response_headers_from_hidden_params, @@ -587,3 +590,55 @@ class TestProviderResponseHeadersInHiddenParams: assert get_provider_response_headers_from_hidden_params(sibling) is None assert "additional_headers" not in sibling._hidden_params + + +def test_spend_metadata_merge_survives_the_logging_thread_inserting_into_the_request_metadata() -> None: + metadata: Final[dict[str, object]] = {} + + def logging_thread_records_hidden_params() -> None: + metadata["hidden_params"] = {"response_cost": 0.1} + + class KeyCheckedWhileTheLoggingThreadRuns(str): + def __contains__(self, key: str) -> bool: + writer: Final = threading.Thread(target=logging_thread_records_hidden_params) + writer.start() + writer.join() + return super().__contains__(key) + + metadata[KeyCheckedWhileTheLoggingThreadRuns("user_api_key_alias")] = "alias" + metadata["user_api_key_hash"] = "hash" + metadata["requester_ip_address"] = "10.0.0.1" + litellm_metadata: Final[dict[str, object]] = {"user_api_key_alias": "stale", "model_group": "claude"} + + merged: Final = get_litellm_metadata_from_kwargs( + {"litellm_params": {"metadata": metadata, "litellm_metadata": litellm_metadata}} + ) + + assert "hidden_params" in metadata + assert merged == {"user_api_key_alias": "alias", "user_api_key_hash": "hash", "model_group": "claude"} + + +def test_spend_metadata_merge_reads_request_metadata_that_is_a_read_only_mapping() -> None: + class ReadOnlyMetadata(Mapping[str, object]): + def __init__(self, values: dict[str, object]) -> None: + self._values: Final = values + + def __getitem__(self, key: str) -> object: + return self._values[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._values) + + def __len__(self) -> int: + return len(self._values) + + merged: Final = get_litellm_metadata_from_kwargs( + { + "litellm_params": { + "metadata": ReadOnlyMetadata({"user_api_key_hash": "hash", "requester_ip_address": "10.0.0.1"}), + "litellm_metadata": {"model_group": "claude"}, + } + } + ) + + assert merged == {"user_api_key_hash": "hash", "model_group": "claude"} diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 02ebae6dcf6..731e842090a 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ( exception_type, extract_and_raise_litellm_exception, ) +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.common_utils import OpenAIError @@ -57,6 +58,104 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st assert caught.value.response.headers["x-request-id"] == "cyber-policy-request" + +def _debug_info_for_api_base(configured_api_base: str) -> str: + request: Final = httpx.Request("POST", "https://gateway.example/v1/chat/completions") + original: Final = OpenAIError( + status_code=500, message="upstream failed", request=request, response=httpx.Response(500, request=request) + ) + with pytest.raises(litellm.InternalServerError) as caught: + exception_type( + model="gpt-5.4-mini", + original_exception=original, + custom_llm_provider="openai", + completion_kwargs={}, + extra_kwargs={"api_base": configured_api_base}, + ) + debug_info: Final = caught.value.litellm_debug_info + assert isinstance(debug_info, str) + return debug_info + + +@pytest.mark.parametrize( + ("configured_api_base", "shown_api_base"), + [ + ("https://qa-user:function-secret@gateway.example/v1?code=function-secret", "https://gateway.example/v1"), + ("https://gateway.example/v1#function-secret", "https://gateway.example/v1"), + ("qa-user:function-secret@gateway.example:8080/v1", "gateway.example:8080/v1"), + ("https://gateway.example?code=function/secret", "https://gateway.example"), + ("https://gateway.example?user=alice@example.org&code=function-secret", "https://gateway.example"), + ("https://gateway.example?email=ops@other-host.example", "https://gateway.example"), + ("qa-user:function-secret@gateway.example/v1?next=https://other.example/y", "gateway.example/v1"), + ( + "https://qa-user:function-secret@gateway.example/models/claude-opus-5-5@default", + "https://gateway.example/models/claude-opus-5-5@default", + ), + ], + ids=[ + "userinfo_and_query", + "fragment", + "no_scheme", + "slash_in_query_without_path", + "at_sign_in_query_without_path", + "at_sign_names_other_host", + "no_scheme_with_url_in_query", + "userinfo_and_at_sign_in_path", + ], +) +def test_api_base_line_keeps_credentials_out(configured_api_base: str, shown_api_base: str) -> None: + debug_info: Final = _debug_info_for_api_base(configured_api_base) + assert f"API Base: `{shown_api_base}`" in debug_info + assert "function-secret" not in debug_info + assert "secret" not in debug_info.split("API Base:", 1)[1].split("\n", 1)[0] + + +@pytest.mark.parametrize( + "configured_api_base", + [ + "https://qa-user:qa?function-secret@gateway.example/v1", + "https://qa-user:qa#function-secret@gateway.example/v1", + "https://qa-user:qa/function-secret@gateway.example/v1", + "https://qa-user:qa/function?secret@gateway.example/v1", + "https://qa.user:qa/function-secret@gateway.example/v1", + "https://qa-user:1234/function-secret@gateway.example/v1", + "https://function-secret/more@gateway.example/v1", + "qa-user:qa://function-secret@gateway.example", + ], + ids=[ + "question_mark_in_password", + "hash_in_password", + "slash_in_password", + "slash_and_question_mark_in_password", + "dotted_username_and_slash_in_password", + "numeric_password_with_slash", + "slash_in_username", + "no_scheme_with_scheme_separator_in_password", + ], +) +def test_api_base_line_is_left_out_when_a_password_breaks_the_url(configured_api_base: str) -> None: + debug_info: Final = _debug_info_for_api_base(configured_api_base) + assert "API Base:" not in debug_info + assert not any(part in debug_info for part in ("qa-user", "qa.user", "function", "secret")) + + +@pytest.mark.parametrize( + "api_base", + [ + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/anthropic/models/claude-opus-5-5@default:rawPredict", + "us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1" + "/publishers/mistralai/models/mistral-large@2407:rawPredict", + "https://gateway.ai.cloudflare.com/v1/account/gateway/workers-ai/@cf/meta/llama-3.1-8b-instruct", + "http://localhost:4000/v1/models/claude@default", + "https://[::1]:8080/v1/models/claude@default", + ], + ids=["vertex_partner_model", "vertex_without_scheme", "workers_ai_path", "localhost", "ipv6"], +) +def test_api_base_line_keeps_an_at_sign_in_the_path(api_base: str) -> None: + assert f"API Base: `{api_base}`" in _debug_info_for_api_base(api_base) + + # Test cases for is_error_str_context_window_exceeded # Tuple format: (error_message, expected_result) context_window_test_cases = [ @@ -480,6 +579,131 @@ def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_lim ) +def test_vertex_status_code_429_ignores_403_token_in_message() -> None: + original_exception: Final = BaseLLMException(status_code=429, message="scripted vertex rejection f4039") + + with pytest.raises(litellm.RateLimitError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 429 + + +def test_vertex_status_code_403_maps_to_403() -> None: + original_exception: Final = BaseLLMException(status_code=403, message="permission denied") + + with pytest.raises(litellm.PermissionDeniedError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_status_code_403_with_text_token_maps_to_bad_request() -> None: + original_exception: Final = BaseLLMException( + status_code=403, + message='{"error":{"code":403,"message":"Permission denied"}}', + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_text_403_without_status_code_still_maps_to_bad_request() -> None: + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=Exception("403 Forbidden: permission denied"), + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_status_substring_without_status_code_does_not_map_to_403() -> None: + with pytest.raises(litellm.APIConnectionError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=Exception("scripted vertex rejection f4039a"), + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code != 403 + + +def test_vertex_wrapped_500_with_403_body_maps_to_bad_request() -> None: + original_exception: Final = BaseLLMException( + status_code=500, + message='Error code: 403 - {"error": {"code": 403, "message": "Permission denied"}}', + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_wrapped_500_with_embedded_403_text_does_not_map_to_403() -> None: + original_exception: Final = BaseLLMException(status_code=500, message="scripted vertex rejection f4039") + + with pytest.raises(litellm.InternalServerError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 500 + + +def test_vertex_wrapped_500_with_body_code_429_and_403_token_maps_to_rate_limit() -> None: + original_exception: Final = BaseLLMException( + status_code=500, + message='{"error":{"code":429,"message":"upstream unavailable (request id: 403-xyz)"}}', + ) + + with pytest.raises(litellm.RateLimitError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 429 + + +def test_vertex_wrapped_500_with_body_code_403_maps_to_bad_request() -> None: + original_exception: Final = BaseLLMException( + status_code=500, + message='{"error":{"code":403,"message":"Permission denied"}}', + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + class TestGetBodyErrorCode: """Unit tests for _get_body_error_code helper.""" @@ -2462,7 +2686,6 @@ def test_openai_gateway_timeout_error(): Test that the OpenAI gateway timeout error is raised """ openai_client = OpenAI() - mapped_target = openai_client.chat.completions.with_raw_response # type: ignore def _return_exception(*args, **kwargs): @@ -2498,8 +2721,8 @@ def test_openai_gateway_timeout_error(): with pytest.raises(litellm.Timeout) as exc_info: with patch.object( - mapped_target, - "create", + openai_client, + "post", side_effect=_return_exception, ): litellm.completion( diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 0256ecf24ed..eca9bbdc170 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -686,6 +686,25 @@ def test_boot_load_does_not_retry_non_retryable_failure(): assert source["fallback_reason"] is not None +def test_boot_load_falls_back_when_urlsplit_rejects_the_url(): + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map( + url="http://[bad/model_prices.json", + sleep=sleeper, + rng=random.Random(0), + client=client, + ) + assert calls["count"] == 0 + assert sleeper.waits == [] + assert _retry_threads() == [] + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert "invalid URL" in (source["fallback_reason"] or "") + + def test_boot_load_respects_local_env_override(monkeypatch): """LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP.""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index dc61d9bc63e..7e23f7988fe 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -7,6 +7,7 @@ import json import logging import os import sys +import threading import time import traceback from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence @@ -39,6 +40,7 @@ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, _get_status_fields, + get_standard_logging_object_payload, set_callbacks, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -51,6 +53,7 @@ from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, + CostBreakdown, ImageResponse, LiteLLMRealtimeStreamLoggingObject, ModelResponse, @@ -7449,6 +7452,75 @@ def test_payload_guardrail_cost_merges_into_existing_cost_breakdown(logging_obj) assert logging_obj.cost_breakdown["total_cost"] == pytest.approx(0.0000429) +def test_cache_hit_payload_keeps_zero_components_beside_guardrail_cost(logging_obj): + logging_obj.set_cost_breakdown( + input_cost=0.00003, + output_cost=0.0000129, + total_cost=0.0000429, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) + payload: Final = _build_success_payload(logging_obj, {**_guardrail_kwargs(response_cost=0.0), "cache_hit": True}) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(0.0003) + assert payload["cost_breakdown"] == { + "input_cost": 0.0, + "output_cost": 0.0, + "tool_usage_cost": 0.0, + "service_tier": None, + "data_residency": None, + "vertex_location": None, + "guardrail_cost": pytest.approx(0.0003), + "total_cost": pytest.approx(0.0003), + } + + +def test_cache_hit_payload_keeps_the_served_pricing_basis_with_every_amount_at_zero(logging_obj): + logging_obj.set_cost_breakdown( + input_cost=0.00003, + output_cost=0.0000129, + total_cost=0.0000474, + cost_for_built_in_tools_cost_usd_dollar=0.0000011, + additional_costs={"azure_model_router_flat_cost": 0.00014}, + original_cost=0.00005, + discount_percent=0.1, + discount_amount=0.000005, + margin_percent=0.05, + margin_fixed_amount=0.000001, + margin_total_amount=0.0000035, + cache_read_cost=0.000002, + cache_creation_cost=0.000003, + reasoning_cost=0.000004, + service_tier="priority", + data_residency="eu", + vertex_location="europe-west4", + ) + kwargs: Final = {"litellm_call_id": "hit-basis-call", "model": "gpt-4o", "messages": [], "cache_hit": True} + + payload: Final = _build_success_payload(logging_obj, kwargs) + + assert payload is not None + assert payload["cost_breakdown"] == { + "input_cost": 0.0, + "output_cost": 0.0, + "total_cost": 0.0, + "tool_usage_cost": 0.0, + "additional_costs": {"azure_model_router_flat_cost": 0.0}, + "original_cost": 0.0, + "discount_percent": 0.1, + "discount_amount": 0.0, + "margin_percent": 0.05, + "margin_fixed_amount": 0.0, + "margin_total_amount": 0.0, + "cache_read_cost": 0.0, + "cache_creation_cost": 0.0, + "reasoning_cost": 0.0, + "service_tier": "priority", + "data_residency": "eu", + "vertex_location": "europe-west4", + } + + def test_payload_without_guardrail_cost_is_unchanged(logging_obj): kwargs = { "litellm_call_id": "no-guardrail-call", @@ -7661,6 +7733,252 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur releasing.async_log_success_event.assert_awaited_once() +def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=stream, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=callbacks, + ) + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) + return logging_obj + + +@pytest.mark.parametrize(("success_callbacks", "expected"), [(["cache"], False), (["cache", "athina"], True)]) +def test_cache_callback_alone_does_not_start_sync_callbacks_for_async_requests( + monkeypatch: pytest.MonkeyPatch, success_callbacks: list[str], expected: bool +) -> None: + monkeypatch.setattr(litellm, "success_callback", success_callbacks) + logging_obj: Final = _async_request_logging_obj("cache-gate", []) + + assert logging_obj._should_run_sync_callbacks_for_async_calls() is expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turn_off_message_logging", [False, True]) +@pytest.mark.parametrize("make_response", [ModelResponse, _responses_api_response_with_text]) +async def test_sync_success_handler_landing_after_logging_hook_keeps_hook_payload_edits( + monkeypatch: pytest.MonkeyPatch, make_response: Callable[[], object], turn_off_message_logging: bool +) -> None: + mirrored: Final = {"x-ratelimit-model_per_project-limit-requests": 600} + response: Final = make_response() + seen_edits: Final[list[tuple[object, object]]] = [] + + class MirroringHook(CustomLogger): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["hidden_params"]["additional_headers"] = dict(mirrored) + payload["response"] = "masked by hook" + logging_obj.success_handler(result=response, for_async_request=True) + return kwargs, result + + class Recorder(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + seen_edits.append((payload["hidden_params"]["additional_headers"], payload["response"])) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "turn_off_message_logging", turn_off_message_logging) + logging_obj: Final = _async_request_logging_obj("late-sync-rebuild", [MirroringHook(), Recorder()]) + + await logging_obj.async_success_handler(result=response) + + assert seen_edits == [(mirrored, "masked by hook")] + + +class _ResponseEditingHook(CustomLogger): + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["response"] = "edited by sync hook" + return kwargs, result + + +def test_sync_logging_hook_edits_stay_off_a_caller_supplied_payload_for_async_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sync_payloads: Final[list[object]] = [] + + class AthinaRecorder: + def log_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + print_verbose: object, + ) -> None: + sync_payloads.append(kwargs["standard_logging_object"]) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", [_ResponseEditingHook(), "athina"]) + monkeypatch.setattr(litellm_logging, "athinaLogger", AthinaRecorder()) + logging_obj: Final = _async_request_logging_obj("caller-supplied-payload", []) + supplied_payload: Final = {"id": "caller-supplied-payload", "response": "original"} + logging_obj.model_call_details["standard_logging_object"] = supplied_payload + + logging_obj.success_handler( + result=ModelResponse(), for_async_request=True, standard_logging_object=supplied_payload + ) + + assert len(sync_payloads) == 1 + sync_payload: Final = sync_payloads[0] + assert isinstance(sync_payload, dict) + assert sync_payload["response"] == "edited by sync hook" + assert supplied_payload == {"id": "caller-supplied-payload", "response": "original"} + assert logging_obj.model_call_details["standard_logging_object"] is supplied_payload + + +@pytest.mark.asyncio +async def test_sync_only_logger_keeps_the_async_payload_when_its_result_builds_none( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sync_payloads: Final[list[object]] = [] + + class AthinaRecorder: + def log_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + print_verbose: object, + ) -> None: + sync_payloads.append(kwargs.get("standard_logging_object")) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", ["athina"]) + monkeypatch.setattr(litellm_logging, "athinaLogger", AthinaRecorder()) + logging_obj: Final = LitellmLogging( + model="pass-through-model", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.datetime.now(), + litellm_call_id="pass-through-sync-logger", + function_id="pass-through-sync-logger", + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + raw_body: Final = json.dumps({"answer": "plain text"}) + await logging_obj.async_success_handler(result=raw_body) + async_payload: Final = logging_obj.model_call_details.get("standard_logging_object") + assert isinstance(async_payload, dict) + async_payload_before_sync: Final = dict(async_payload) + monkeypatch.setattr(litellm, "success_callback", [_ResponseEditingHook(), "athina"]) + + logging_obj.success_handler(result=raw_body, for_async_request=True) + + assert sync_payloads == [{**async_payload_before_sync, "response": "edited by sync hook"}] + assert logging_obj.model_call_details["standard_logging_object"] is async_payload + assert async_payload == async_payload_before_sync + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("entrypoint", ["handle_sync_success_callbacks_for_async_calls", "dispatch_success_handlers"]) +@pytest.mark.parametrize( + ("callback_name", "logger_global"), + [ + ("athina", "athinaLogger"), + ("promptlayer", "promptLayerLogger"), + ("wandb", "weightsBiasesLogger"), + ("lunary", "lunaryLogger"), + ("helicone", "heliconeLogger"), + ("greenscale", "greenscaleLogger"), + ("s3", "s3Logger"), + ("traceloop", "traceloopLogger"), + ], +) +async def test_sync_only_logger_gets_its_own_payload_for_async_requests( + monkeypatch: pytest.MonkeyPatch, stream: bool, entrypoint: str, callback_name: str, logger_global: str +) -> None: + sync_logged: Final = threading.Event() + sync_payloads: Final[list[object]] = [] + async_payloads: Final[list[object]] = [] + + class LegacyLoggerRecorder: + def _record(self, **call: object) -> None: + kwargs = call["kwargs"] + assert isinstance(kwargs, dict) + sync_payloads.append(kwargs["standard_logging_object"]) + sync_logged.set() + + log_event = log_success = _record + + class AsyncRecorder(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + async_payloads.append(kwargs["standard_logging_object"]) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", [callback_name]) + monkeypatch.setattr(litellm_logging, logger_global, LegacyLoggerRecorder(), raising=False) + response: Final = ModelResponse(id="chatcmpl-sync-only") + logging_obj: Final = _async_request_logging_obj("sync-only-logger", [AsyncRecorder()], stream=stream) + now: Final = datetime.datetime.now() + + if entrypoint == "dispatch_success_handlers": + await logging_obj.dispatch_success_handlers(result=response, start_time=now, end_time=now) + else: + await logging_obj.async_success_handler(result=response, start_time=now, end_time=now) + logging_obj.handle_sync_success_callbacks_for_async_calls(result=response, start_time=now, end_time=now) + + assert sync_logged.wait(timeout=10) + assert len(async_payloads) == 1 + assert len(sync_payloads) == 1 + async_payload: Final = async_payloads[0] + sync_payload: Final = sync_payloads[0] + assert isinstance(async_payload, dict) + assert isinstance(sync_payload, dict) + assert sync_payload["id"] == async_payload["id"] + assert sync_payload is not async_payload + assert logging_obj.model_call_details["standard_logging_object"] is async_payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_supabase_logs_the_request_messages_and_response_for_async_requests( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + supabase_logged: Final = threading.Event() + supabase_calls: Final[list[tuple[object, object]]] = [] + + class SupabaseRecorder: + def log_event(self, **call: object) -> None: + supabase_calls.append((call["messages"], call["response_obj"])) + supabase_logged.set() + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", ["supabase"]) + monkeypatch.setattr(litellm_logging, "supabaseClient", SupabaseRecorder()) + response: Final = ModelResponse(id="chatcmpl-supabase") + logging_obj: Final = _async_request_logging_obj("supabase-logger", [], stream=stream) + now: Final = datetime.datetime.now() + + await logging_obj.dispatch_success_handlers(result=response, start_time=now, end_time=now) + + assert supabase_logged.wait(timeout=10) + assert len(supabase_calls) == 1 + messages, response_obj = supabase_calls[0] + assert messages == [{"role": "user", "content": "hi"}] + assert isinstance(response_obj, ModelResponse) + assert response_obj.id == "chatcmpl-supabase" + + def _resolve(custom_llm_provider, litellm_params, optional_params, model): from litellm.litellm_core_utils.litellm_logging import ( _resolve_vertex_location_for_cost, @@ -9990,6 +10308,245 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch): is_otel_v2_enabled.cache_clear() +def _lit8679_priced_logging_obj(deployment_id: str, litellm_call_id: str) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-5.4-nano", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=time.time(), + litellm_call_id=litellm_call_id, + function_id="fn", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + user="", + optional_params={}, + litellm_params={"metadata": {"model_group": "lit8679-group", "model_info": {"id": deployment_id}}}, + custom_llm_provider="openai", + ) + return logging_obj + + +def test_uncached_price_recomputes_price_past_a_zero_stamp(): + """A probe for the price a cache-served response would have had ignores the 0.0 stamped on it.""" + deployment_id: Final = "lit8679-uncached-cost-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-uncached") + oracle_logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-oracle") + cached_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + cached_response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + fresh_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + cost: Final = logging_obj.uncached_price(result=cached_response).cost + uncached_price: Final = oracle_logging_obj._response_cost_calculator(result=fresh_response) + + assert cost == uncached_price + assert cost is not None and cost > 0.0 + assert logging_obj.cost_breakdown is None + finally: + litellm.model_cost.pop(deployment_id, None) + + +def test_uncached_price_leaves_no_pricing_state_on_the_call(): + """The uncached price probe leaves no cost breakdown or billed rates behind on the cache-hit call.""" + deployment_id: Final = "lit8679-probe-state-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-probe-state") + cached_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + cached_response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + assert (logging_obj.cost_breakdown, logging_obj.billed_token_rates) == (None, None) + + cost: Final = logging_obj.uncached_price(result=cached_response).cost + + assert cost is not None and cost > 0.0 + assert (logging_obj.cost_breakdown, logging_obj.billed_token_rates) == (None, None) + finally: + litellm.model_cost.pop(deployment_id, None) + + +def _assert_zeroed_copy(hit_breakdown: CostBreakdown | None, fresh_breakdown: CostBreakdown) -> None: + basis_and_rates: Final = ("service_tier", "data_residency", "vertex_location", "discount_percent", "margin_percent") + assert hit_breakdown is not None and hit_breakdown.keys() == fresh_breakdown.keys() + assert all(hit_breakdown[key] == fresh_breakdown[key] for key in basis_and_rates if key in hit_breakdown) + assert all(value == 0.0 for key, value in hit_breakdown.items() if key not in basis_and_rates) + + +def test_payload_saved_cache_cost_keeps_uncached_price_when_result_stamped_zero(): + """A cache-hit payload still reports what the call would have cost, not the 0.0 stamped on the served response.""" + deployment_id: Final = "lit8679-saved-cost-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-payload") + oracle_logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-payload-oracle") + response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + fresh_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + logging_obj.model_call_details["cache_hit"] = True + now: Final = datetime.datetime.now() + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-5.4-nano", + "messages": [], + "cache_hit": True, + "litellm_params": logging_obj.litellm_params, + }, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + uncached_price: Final = oracle_logging_obj._response_cost_calculator(result=fresh_response) + + assert payload is not None + assert payload["saved_cache_cost"] == uncached_price + assert payload["saved_cache_cost"] > 0.0 + assert payload["response_cost"] == 0.0 + assert logging_obj.cost_breakdown is None + assert oracle_logging_obj.cost_breakdown is not None + _assert_zeroed_copy(payload["cost_breakdown"], oracle_logging_obj.cost_breakdown) + finally: + litellm.model_cost.pop(deployment_id, None) + + +def test_cache_hit_payload_reports_zero_cost_though_the_call_priced_the_response_as_fresh(): + """The proxy prices a /v1/messages body for its cost header before the success callback learns it was a cache + hit, and a replayed stream's assembled response carries its original price, so the call holds full-price numbers + that the hit's payload must not export.""" + deployment_id: Final = "lit8679-stale-breakdown-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-stale-breakdown") + response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + header_price: Final = logging_obj.response_cost_calculator(result=response) + response._hidden_params["response_cost"] = header_price + logging_obj.model_call_details["cache_hit"] = True + now: Final = datetime.datetime.now() + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-5.4-nano", + "messages": [], + "cache_hit": True, + "response_cost": 0.0, + "litellm_params": logging_obj.litellm_params, + }, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert header_price is not None and header_price > 0.0 + assert logging_obj.cost_breakdown is not None + assert payload is not None + assert payload["response_cost"] == 0.0 + assert payload["hidden_params"]["response_cost"] == 0.0 + _assert_zeroed_copy(payload["cost_breakdown"], logging_obj.cost_breakdown) + assert payload["saved_cache_cost"] == header_price + finally: + litellm.model_cost.pop(deployment_id, None) + + +def test_uncached_price_leaves_no_failure_debug_info_for_unpriced_model(): + """A probe that cannot price the model returns None without stamping the pricing failure on the cache-hit call.""" + logging_obj: Final = LitellmLogging( + model="lit8679-model-with-no-price", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=time.time(), + litellm_call_id="lit8679-unpriced-probe", + function_id="fn", + ) + logging_obj.update_environment_variables( + model="lit8679-model-with-no-price", + user="", + optional_params={}, + litellm_params={"metadata": {}}, + custom_llm_provider="openai", + ) + cached_response: Final = ModelResponse( + model="lit8679-model-with-no-price", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + cached_response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + assert "response_cost_failure_debug_information" not in logging_obj.model_call_details + + cost: Final = logging_obj.uncached_price(result=cached_response).cost + + assert cost is None + assert "response_cost_failure_debug_information" not in logging_obj.model_call_details + + def test_enterprise_alerting_loggers_resolves_real_classes(): from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 9b2c42450ae..7d85da3f756 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -13,6 +13,7 @@ from websockets.frames import Close import litellm from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -543,6 +544,32 @@ async def test_backend_to_client_stores_normalized_events_for_logging(): assert streaming.messages[0]["response"]["usage"]["total_tokens"] == 0 +@pytest.mark.asyncio +async def test_client_ack_messages_replays_the_opening_message_a_refused_bedrock_stream_left(): + opening_message: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is 17 * 3?"}]}, + } + ) + next_message: Final = json.dumps({"type": "response.create"}) + client_ws = MagicMock() + client_ws.scope = {"headers": [], BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: opening_message} + client_ws.receive_text = AsyncMock(side_effect=[next_message, Exception("connection closed")]) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in backend_ws.send.call_args_list] + assert sent_types == ["conversation.item.create", "response.create"] + assert json.loads(backend_ws.send.call_args_list[0].args[0])["item"]["content"][0]["text"] == "What is 17 * 3?" + assert BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY not in client_ws.scope + + @pytest.mark.asyncio async def test_client_ack_messages_keeps_beta_session_shape_for_beta_clients(): client_ws = MagicMock() diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 93f700b1c2a..c07fe2d1c31 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2,14 +2,17 @@ import asyncio import json import threading import time -from typing import Final, NoReturn, Optional +from typing import AsyncIterator, Final, Iterator, Literal, NoReturn, Optional from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import pytest import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( CustomStreamWrapper, @@ -1382,6 +1385,140 @@ async def test_async_streaming_read_timeout_triggers_midstream_fallback( assert isinstance(excinfo.value.original_exception, Exception) +_AZURE_DEPLOYMENT_API_BASE: Final = "https://mid-stream-debug.openai.azure.com" +_AZURE_DEPLOYMENT_REQUEST_URL: Final = ( + f"{_AZURE_DEPLOYMENT_API_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2025-04-01-preview" +) +_MID_STREAM_START_TIME: Final = 1_790_000_000.0 + + +def _azure_deployment_logging_obj(logged_api_base: str | None, routed_api_base: str | None) -> Logging: + logging_obj: Final = Logging( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "Count from 1 to 40"}], + stream=True, + call_type="acompletion", + start_time=_MID_STREAM_START_TIME, + litellm_call_id="mid-stream-debug-info", + function_id="mid-stream-debug-info", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-mini", + optional_params={}, + litellm_params=get_litellm_params( + api_base=logged_api_base, + custom_llm_provider="azure", + metadata={ + "model_group": "gpt-5.4-mini-azure", + "deployment": "azure/gpt-5.4-mini", + "api_base": routed_api_base, + }, + ), + ) + return logging_obj + + +def _content_then_failure(failure: Literal["transport_error", "unparseable_chunk"]) -> Iterator[object]: + yield ModelResponseStream( + model="gpt-5.4-mini-2026-03-17", + choices=[StreamingChoices(index=0, delta=Delta(content="1", role="assistant"))], + ) + if failure == "transport_error": + raise httpx.ReadError("Response payload is not completed") + yield object() + + +async def _as_async(chunks: Iterator[object]) -> AsyncIterator[object]: + for chunk in chunks: + yield chunk + + +async def _drain(wrapper: CustomStreamWrapper, sync_mode: bool) -> None: + if sync_mode: + for _ in wrapper: + pass + return + async for _ in wrapper: + pass + + +async def _mid_stream_debug_info( + logged_api_base: str | None, + routed_api_base: str | None, + failure: Literal["transport_error", "unparseable_chunk"], + sync_mode: bool, +) -> str: + chunks: Final = _content_then_failure(failure) + wrapper: Final = CustomStreamWrapper( + completion_stream=chunks if sync_mode else _as_async(chunks), + model="gpt-5.4-mini-2026-03-17", + logging_obj=_azure_deployment_logging_obj(logged_api_base, routed_api_base), + custom_llm_provider="azure", + ) + with pytest.raises(MidStreamFallbackError) as excinfo: + await _drain(wrapper, sync_mode) + assert excinfo.value.is_pre_first_chunk is False + debug_info: Final = getattr(excinfo.value.original_exception, "litellm_debug_info", None) + assert isinstance(debug_info, str) + return debug_info + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("failure", ["transport_error", "unparseable_chunk"]) +@pytest.mark.parametrize( + ("logged_api_base", "routed_api_base"), + [ + (_AZURE_DEPLOYMENT_REQUEST_URL, _AZURE_DEPLOYMENT_API_BASE), + (_AZURE_DEPLOYMENT_API_BASE, None), + (_AZURE_DEPLOYMENT_API_BASE, ""), + ], + ids=["routed", "sdk", "empty_routed"], +) +async def test_mid_stream_error_reports_the_deployment_it_was_sent_to( + logged_api_base: str, + routed_api_base: str | None, + failure: Literal["transport_error", "unparseable_chunk"], + sync_mode: bool, +): + debug_info: Final = await _mid_stream_debug_info(logged_api_base, routed_api_base, failure, sync_mode) + + assert f"API Base: `{_AZURE_DEPLOYMENT_API_BASE}`" in debug_info + assert "model_group: `gpt-5.4-mini-azure`" in debug_info + assert "deployment: `azure/gpt-5.4-mini`" in debug_info + assert "Count from 1 to 40" not in debug_info + assert "api.openai.com" not in debug_info + + +@pytest.mark.asyncio +async def test_mid_stream_error_leaves_out_an_api_base_it_cannot_resolve(): + debug_info: Final = await _mid_stream_debug_info(None, None, "transport_error", sync_mode=False) + + assert "model_group: `gpt-5.4-mini-azure`" in debug_info + assert "API Base:" not in debug_info + + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routed_api_base", + [ + f"{_AZURE_DEPLOYMENT_API_BASE}?subscription-key=deployment-secret-value", + f"{_AZURE_DEPLOYMENT_API_BASE}?code=deployment-secret-value", + _AZURE_DEPLOYMENT_API_BASE.replace("https://", "https://qa-user:deployment-secret-value@"), + ], + ids=["subscription_key", "function_code", "userinfo"], +) +async def test_mid_stream_error_keeps_credentials_in_the_api_base_out(routed_api_base: str): + debug_info: Final = await _mid_stream_debug_info( + _AZURE_DEPLOYMENT_REQUEST_URL, routed_api_base, "transport_error", sync_mode=False + ) + + assert f"API Base: `{_AZURE_DEPLOYMENT_API_BASE}`" in debug_info + assert "deployment-secret" not in debug_info + + + def test_streaming_handler_with_created_time_propagation( initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging ): diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 892736b5043..0aca60351bf 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -6,7 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic """ import json -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch import pytest @@ -18,7 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import Streami from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs class MockPassThroughGuardrail(CustomGuardrail): @@ -383,6 +383,97 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert '"name": "lookup_fruit"' not in raw assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + @staticmethod + def _message_delta(chunks: list[bytes]) -> dict[str, object]: + lines: Final = b"".join(chunks).decode().split("\n") + events: Final = (json.loads(line[len("data:") :]) for line in lines if line.startswith("data:")) + return next(event["delta"] for event in events if event.get("type") == "message_delta") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_text_rewrite_drops_the_upstream_explanations(self): + delta_fields: Final = { + "stop_details": {"type": "refusal", "category": "cyber", "explanation": "the reply says world"}, + "safeguard_results": [ + {"type": "prompt_injection", "status": {"type": "available", "explanation": "quotes world"}} + ], + } + source: Final[list[bytes]] = self._ended_sse_chunks() + chunks: Final = [ + chunk.replace( + b'"stop_sequence": null}', b'"stop_sequence": null, ' + json.dumps(delta_fields)[1:-1].encode() + b"}" + ) + for chunk in source + ] + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._delta_texts(chunks) == ["hello [MASKED]", ""] + assert self._message_delta(chunks) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": None}, + "safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}], + } + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_drops_the_verdict_of_the_rewritten_call_only(self): + class MaskPersimmon(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + for tool_call in inputs.get("tool_calls", []): + if isinstance(tool_call, ChatCompletionMessageToolCall) and "persimmon" in tool_call.function.arguments: + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + first_call: Final[list[bytes]] = self._ended_tool_use_sse_chunks() + second_call: Final = [ + chunk.replace(b'"index": 0', b'"index": 1') + .replace(b"toolu_1", b"toolu_2") + .replace(b"persim", b"pear") + .replace(b'mon\\"}', b'\\"}') + for chunk in first_call[1:6] + ] + verdicts: Final = { + "toolu_1": {"type": "evaluated", "outcome": "flagged", "explanation": "looks up persimmon"}, + "toolu_2": {"type": "evaluated", "outcome": "not_flagged", "explanation": "looks up pear"}, + } + safeguard_results: Final = [ + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": verdicts}} + ] + message_delta: Final = first_call[6].replace( + b'"stop_sequence": null}', + b'"stop_sequence": null, "safeguard_results": ' + json.dumps(safeguard_results).encode() + b"}", + ) + chunks: Final = [*first_call[:6], *second_call, message_delta, first_call[7]] + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MaskPersimmon(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._partial_jsons(chunks)[:3] == ['{"fruit": "[MASKED]"}', "", ""] + assert self._message_delta(chunks)["safeguard_results"] == [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"toolu_2": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + @pytest.mark.asyncio async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 724e137c835..6323132db36 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1346,6 +1346,74 @@ def test_opus_uses_native_structured_output(model_name): assert optional_params.get("json_mode") is True +@pytest.mark.parametrize( + ("field_schema", "native"), + [ + ({"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, True), + ({"type": "object", "additionalProperties": {"type": "integer"}}, False), + ({"anyOf": [{"type": "object"}, {"type": "null"}]}, False), + ], + ids=["closed-object", "dict-field", "nullable-dict-field"], +) +def test_dict_fields_take_the_json_tool_on_a_native_structured_output_model(field_schema, native): + optional_params: Final = AnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "out", + "schema": {"type": "object", "properties": {"field": field_schema}, "required": ["field"]}, + }, + } + }, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + + assert ("output_format" in optional_params) is native + assert ([tool["name"] for tool in optional_params.get("tools", [])] == ["json_tool_call"]) is not native + + + +@pytest.mark.parametrize( + ("model", "thinking_params"), + [ + ("claude-opus-4-8", {"thinking": {"type": "adaptive"}}), + ("claude-opus-4-8", {"reasoning_effort": "low"}), + ("claude-opus-4-8", {"reasoning_effort": "none"}), + ("claude-sonnet-5-5", {}), + ( + "claude-opus-4-8", + {"tools": [{"type": "function", "function": {"name": "lookup", "parameters": {"type": "object"}}}]}, + ), + ], + ids=["adaptive-thinking", "reasoning-effort", "reasoning-effort-none", "model-without-forced-tools", "caller-tools"], +) +def test_dict_fields_keep_native_output_when_the_json_tool_would_not_be_forced(model, thinking_params): + optional_params: Final = AnthropicConfig().map_openai_params( + non_default_params={ + **thinking_params, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "out", + "schema": { + "type": "object", + "properties": {"field": {"type": "object", "additionalProperties": {"type": "integer"}}}, + "required": ["field"], + }, + }, + }, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "output_format" in optional_params + assert "json_tool_call" not in [tool["name"] for tool in optional_params.get("tools", [])] + def test_native_structured_output_uses_bundled_capability_when_remote_map_lags( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -9377,17 +9445,26 @@ def test_map_openai_params_lets_a_native_anthropic_search_tool_replace_the_opena @pytest.mark.parametrize( - ("config", "model", "provider"), + ("config", "model", "provider", "model_cost_key"), [ - pytest.param(AmazonAnthropicClaudeConfig(), "us.anthropic.claude-sonnet-5-5", "bedrock", id="bedrock-invoke"), - pytest.param(VertexAIAnthropicConfig(), "claude-sonnet-5-5", "vertex_ai", id="vertex_ai"), + pytest.param( + AmazonAnthropicClaudeConfig(), + "us.anthropic.claude-sonnet-5-5", + "bedrock", + "us.anthropic.claude-sonnet-5-5", + id="bedrock-invoke", + ), + pytest.param( + VertexAIAnthropicConfig(), "claude-sonnet-5-5", "vertex_ai", "vertex_ai/claude-sonnet-5-5", id="vertex_ai" + ), ], ) def test_map_openai_params_judges_tool_search_on_the_real_model_when_tool_based_response_format_swaps_it( - config, model, provider + set_native_structured_output, config, model, provider, model_cost_key ): from litellm.responses.litellm_completion_transformation.tool_search import anthropic_tool_search_chat_tool + set_native_structured_output(model_cost_key, None) assert litellm.supports_tool_search(model, provider) mapped: Final = config.map_openai_params( non_default_params={ @@ -9549,6 +9626,17 @@ def test_json_mode_tool_use_turn_adds_no_stop_fields_to_the_message(): assert not message.provider_specific_fields +def test_json_mode_keeps_the_thinking_that_preceded_the_json_tool_call(): + thinking_block: Final = {"type": "thinking", "thinking": "The user wants a greeting.", "signature": "sig-1"} + body: Final = _json_mode_message_body("tool_use", None) + result: Final = _transform_json_mode_body({**body, "content": [thinking_block, *body["content"]]}) + + message: Final = result.choices[0].message + assert json.loads(message.content) == {"answer": "Hello"} + assert message.reasoning_content == thinking_block["thinking"] + assert [block["signature"] for block in message.thinking_blocks] == [thinking_block["signature"]] + + def test_non_dict_upstream_body_still_transforms_without_refusal_fields() -> None: result = AnthropicConfig().transform_parsed_response( completion_response=MagicMock(), diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py index b73df255cad..7400cd08dd3 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -18,6 +18,10 @@ import anthropic import httpx import pytest +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, _CombinedChunkSplitter, @@ -143,6 +147,43 @@ def test_delayed_usage_chunk_preserves_cache_tokens(): assert message_delta["usage"]["cache_creation_input_tokens"] == 20 +def _cached_bedrock_response(cache_read: int, cache_write: int) -> dict: + return { + "id": "chatcmpl-cached", + "object": "chat.completion", + "created": 1700000000, + "model": "us.anthropic.claude-opus-4-7", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}], + "usage": Usage( + prompt_tokens=42145, + completion_tokens=6, + total_tokens=42151, + cache_read_input_tokens=cache_read, + cache_creation_input_tokens=cache_write, + ).model_dump(), + } + + +async def _replay_async(cached: dict) -> list[ModelResponseStream]: + return [chunk async for chunk in convert_to_streaming_response_async(cached)] + + +@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +def test_a_cache_hit_replayed_as_a_messages_stream_keeps_the_prompt_cache_tokens(replay, cache_read, cache_write): + cached: Final = _cached_bedrock_response(cache_read, cache_write) + chunks: Final = ( + asyncio.run(_replay_async(cached)) if replay == "async" else list(convert_to_streaming_response(cached)) + ) + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-opus-4-7")) + message_delta: Final = next(event for event in events if event.get("type") == "message_delta") + + assert message_delta["usage"]["input_tokens"] == 6 + assert message_delta["usage"].get("cache_read_input_tokens", 0) == cache_read + assert message_delta["usage"].get("cache_creation_input_tokens", 0) == cache_write + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)]) diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py new file mode 100644 index 00000000000..476cea2833c --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py @@ -0,0 +1,93 @@ +import json +from collections.abc import Iterable, Mapping +from typing import Final + +from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, +) + +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] + + +def _decoded_event(chunk: object) -> tuple[str, dict[str, object]]: + assert isinstance(chunk, bytes) + event_line, data_line = chunk.decode().split("\n")[:2] + assert event_line.startswith("event: ") + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +def _decoded_events(chunks: Iterable[object]) -> tuple[tuple[str, dict[str, object]], ...]: + return tuple(_decoded_event(chunk) for chunk in chunks) + + +def _message_delta(response: Mapping[str, object]) -> object: + decoded: Final = _decoded_events(FakeAnthropicMessagesStreamIterator(response=response)) + assert tuple(event_type for event_type, _ in decoded)[-2:] == ("message_delta", "message_stop") + deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta") + assert len(deltas) == 1 + return deltas[0] + + +def test_message_delta_carries_safeguard_results(): + """Websearch interception and the MCP loop render a complete response as + this stream for a client that asked to stream (the guardrail path does too, + after merging the upstream fields back in, which + tests/unit/proxy/guardrails/test_anthropic_sse.py covers); a Claude Code + auto-mode client reads the classifier verdict from message_delta.delta, + where a real Anthropic stream puts it.""" + response: Final = { + "id": "msg_1", + "model": "claude-sonnet-4-6", + "role": "assistant", + "type": "message", + "content": [{"type": "text", "text": "I won't run that."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + assert _message_delta(response) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + + +def test_message_delta_carries_refusal_stop_details(): + stop_details: Final = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response: Final = { + "id": "msg_2", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "usage": {"input_tokens": 10, "output_tokens": 0}, + } + assert _message_delta(response) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + + +def test_message_delta_omits_absent_optional_fields(): + response: Final = { + "id": "msg_3", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": {"input_tokens": 10, "output_tokens": 2}, + } + assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py index 7e983be835b..b995c6dba9b 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py @@ -1306,6 +1306,33 @@ def test_anthropic_messages_streaming_response_reports_withheld_output_of_its_st ) +def test_anthropic_messages_response_as_sse_events_keeps_message_delta_fields_out_of_message_start(): + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}] + stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response = { + "id": "msg_1", + "model": "claude-haiku", + "role": "assistant", + "type": "message", + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "safeguard_results": safeguard_results, + "content": [{"type": "text", "text": "no"}], + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + message_start = decoded[0][1]["message"] + assert message_start["stop_details"] is None + assert "safeguard_results" not in message_start + assert decoded[-2][1]["delta"] == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "safeguard_results": safeguard_results, + } + + def test_anthropic_messages_response_as_sse_events_no_content_blocks(): response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"} decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) @@ -1530,6 +1557,72 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co assert deferred_fired.is_set() +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] + + +def _sse_events_message_delta(response: Mapping[str, object]) -> object: + decoded: Final = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta") + assert len(deltas) == 1 + return deltas[0] + + +def test_anthropic_messages_response_as_sse_events_message_delta_carries_safeguard_results(): + """A Claude Code auto-mode request's classifier verdict rides in + message_delta.delta on a real Anthropic stream; a synthesized stream that + drops it leaves the client without the verdict.""" + response = { + "id": "msg_7", + "content": [{"type": "text", "text": "I won't run that."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + assert _sse_events_message_delta(response) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + + +def test_anthropic_messages_response_as_sse_events_message_delta_carries_refusal_stop_details(): + stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response = { + "id": "msg_8", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + assert _sse_events_message_delta(response) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + + +def test_anthropic_messages_response_as_sse_events_message_delta_omits_absent_optional_fields(): + response = { + "id": "msg_9", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + } + assert _sse_events_message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} + + class _OneVerdictClassifier: def __init__(self, reply: str) -> None: self._reply = reply diff --git a/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py b/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py index 33f95ccc9ff..dbd49064e4b 100644 --- a/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py @@ -4,6 +4,7 @@ from typing import Final import pytest from litellm.llms.anthropic.pass_through.stream_assembly import ( + anthropic_message_delta_fields, assemble_anthropic_sse_stream, is_anthropic_sse_stream, is_raw_sse_stream, @@ -55,6 +56,35 @@ ANTHROPIC_STREAM: Final = ( _frame("message_stop", {"type": "message_stop"}), ) +SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] +SAFEGUARDED_STREAM: Final = ( + *ANTHROPIC_STREAM[:5], + _frame( + "message_delta", + { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": SAFEGUARD_RESULTS, + }, + "usage": {"output_tokens": 2}, + }, + ), + _frame("message_stop", {"type": "message_stop"}), +) + GOOGLE_STREAM: Final = (b'data: {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}\n\n',) ANTHROPIC_ERROR: Final = _frame("error", {"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}) CHAT_ERROR: Final = b'data: {"error": {"message": "blocked by guardrail"}}\n\n' @@ -108,6 +138,12 @@ def test_assemble_anthropic_sse_stream_restores_upstream_identity_only_when_aske assert (restored.id, restored.model) == (MESSAGE_ID, MODEL) +def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops(): + assert dict(anthropic_message_delta_fields(ANTHROPIC_STREAM)) == {} + assert dict(anthropic_message_delta_fields(SAFEGUARDED_STREAM)) == {"safeguard_results": SAFEGUARD_RESULTS} + assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {} + + @pytest.mark.parametrize( "chunks", [GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()], diff --git a/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py index b2850192f91..0ed3b907a6f 100644 --- a/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py +++ b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py @@ -12,6 +12,18 @@ from litellm.llms.anthropic.chat.transformation import AnthropicConfig class TestOutputFormatArrayObjectConstraints: + def test_closes_an_object_whose_type_is_a_union_with_null(self): + schema = { + "type": "object", + "properties": { + "address": {"type": ["object", "null"], "properties": {"city": {"type": "string"}}}, + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["properties"]["address"]["additionalProperties"] is False + def test_removes_uniqueitems(self): schema = { "type": "object", @@ -129,7 +141,7 @@ class TestOutputFormatArrayObjectConstraints: schema = { "type": "array", "items": {"type": "string"}, - "minItems": 1, + "minItems": 2, "maxItems": 10, "uniqueItems": True, "minContains": 2, @@ -139,7 +151,7 @@ class TestOutputFormatArrayObjectConstraints: result = AnthropicConfig.filter_anthropic_output_schema(schema) assert result["description"] == ( - "Note: minimum number of items: 1, maximum number of items: 10, " + "Note: minimum number of items: 2, maximum number of items: 10, " "all array items must be unique, minimum number of matching items: 2, " "maximum number of matching items: 3." ) diff --git a/tests/unit/llms/anthropic/test_anthropic_schema_filter.py b/tests/unit/llms/anthropic/test_anthropic_schema_filter.py index 32423d6679f..407a52c183e 100644 --- a/tests/unit/llms/anthropic/test_anthropic_schema_filter.py +++ b/tests/unit/llms/anthropic/test_anthropic_schema_filter.py @@ -69,7 +69,7 @@ class TestFilterAnthropicOutputSchema: schema = { "type": "array", "items": {"type": "string"}, - "minItems": 1, + "minItems": 2, "maxItems": 10, } @@ -80,9 +80,27 @@ class TestFilterAnthropicOutputSchema: assert result["type"] == "array" assert result["items"] == {"type": "string"} # Description should contain constraint info - assert "minimum number of items: 1" in result["description"] + assert "minimum number of items: 2" in result["description"] assert "maximum number of items: 10" in result["description"] + def test_keeps_min_items_of_zero_or_one(self): + for min_items in (0, 1): + schema = {"type": "array", "items": {"type": "string"}, "minItems": min_items} + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["minItems"] == min_items + assert "description" not in result + + def test_moves_a_min_items_that_is_not_an_integer_into_the_description(self): + for min_items in (True, 1.0): + schema = {"type": "array", "items": {"type": "string"}, "minItems": min_items} + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "minItems" not in result + assert "minimum number of items" in result["description"] + def test_handles_nested_schemas(self): """Test that nested schemas are also filtered.""" schema = { @@ -100,7 +118,7 @@ class TestFilterAnthropicOutputSchema: } }, }, - "minItems": 1, + "minItems": 2, } }, } @@ -336,7 +354,7 @@ class TestFilterAnthropicOutputSchema: schema = { "type": "array", "items": {"type": "string"}, - "minItems": 1, + "minItems": 2, "maxItems": 10, "uniqueItems": True, "minContains": 2, @@ -346,7 +364,7 @@ class TestFilterAnthropicOutputSchema: result = AnthropicConfig.filter_anthropic_output_schema(schema) assert result["description"] == ( - "Note: minimum number of items: 1, maximum number of items: 10, " + "Note: minimum number of items: 2, maximum number of items: 10, " "all array items must be unique, minimum number of matching items: 2, " "maximum number of matching items: 3." ) diff --git a/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py index 3630ef1e354..1cdc288ff87 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -1,9 +1,10 @@ import json import traceback -from typing import Callable, Optional +from typing import Callable, Final, Optional from unittest.mock import MagicMock, patch import pytest +from openai._legacy_response import LegacyAPIResponse import litellm from litellm import ModelResponse @@ -130,7 +131,7 @@ def test_azure_o3_streaming(): api_version="2024-02-15-preview", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: try: completion( model="azure/o3-mini", @@ -141,7 +142,7 @@ def test_azure_o3_streaming(): except Exception as e: print(e) assert mock_create.call_count == 1 - assert "stream" in mock_create.call_args.kwargs + assert "stream" in mock_create.call_args.kwargs["body"] def test_azure_o_series_routing(): @@ -157,7 +158,7 @@ def test_azure_o_series_routing(): api_version="2024-02-15-preview", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: try: completion( model="azure/o_series/my-random-deployment-name", @@ -168,17 +169,19 @@ def test_azure_o_series_routing(): except Exception as e: print(e) assert mock_create.call_count == 1 - assert "stream" not in mock_create.call_args.kwargs + assert "stream" not in mock_create.call_args.kwargs["body"] @patch("litellm.main.azure_o1_chat_completions._get_openai_client") def test_openai_o_series_max_retries_0(mock_get_openai_client): import litellm - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( - ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) - ) + raw_response: Final = MagicMock(spec=LegacyAPIResponse) + raw_response.headers = {} + raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + client_stub: Final = MagicMock() + client_stub.post.return_value = raw_response + mock_get_openai_client.return_value = client_stub litellm.set_verbose = True response = litellm.completion( model="azure/o1-preview", @@ -233,7 +236,7 @@ async def test_azure_o1_series_response_format_extra_params(): response_format = {"type": "json_object"} tool_choice = "auto" with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -251,7 +254,7 @@ async def test_azure_o1_series_response_format_extra_params(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", json.dumps(request_body, indent=4)) assert request_body["tools"] == tools diff --git a/tests/unit/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py index 7907344ae64..7d5afde68ef 100644 --- a/tests/unit/llms/azure/test_azure.py +++ b/tests/unit/llms/azure/test_azure.py @@ -313,7 +313,7 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): number: str = Field(description="total number of days in a week") days: list[str] = Field(description="name of days in a week") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + with patch.object(client, "post") as mock_post: try: completion( model="azure/gpt4o", @@ -326,9 +326,9 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): mock_post.assert_called_once() - print(mock_post.call_args.kwargs) + print(mock_post.call_args.kwargs["body"]) - assert "tool_calls" not in mock_post.call_args.kwargs + assert "tool_calls" not in mock_post.call_args.kwargs["body"] @pytest.mark.parametrize( @@ -382,7 +382,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): } ] - with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + with patch.object(client, "post") as mock_post: mock_post.return_value.headers = {} mock_post.return_value.parse.return_value = litellm.ModelResponse( choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}] @@ -411,9 +411,9 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): mock_post.assert_called_once() if api_version == "2024-10-21": - assert "response_format" in mock_post.call_args.kwargs + assert "response_format" in mock_post.call_args.kwargs["body"] else: - assert "response_format" not in mock_post.call_args.kwargs + assert "response_format" not in mock_post.call_args.kwargs["body"] assert response.choices[0].message.content == InvestigationOutput().model_dump_json() diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index a445d951bc2..634228da1bf 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -384,6 +384,66 @@ def test_azure_ai_grok_bills_reasoning_tokens_reported_outside_completion_tokens assert response._hidden_params["response_cost"] == pytest.approx(expected_cost, rel=1e-9) +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("upstream_model", "usage", "expected_completion_tokens"), + [ + pytest.param( + "grok-4.7", + { + "prompt_tokens": 14, + "completion_tokens": 10, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + 322, + id="grok-reports-reasoning-outside-completion", + ), + pytest.param( + "gpt-5.4-mini", + { + "prompt_tokens": 14, + "completion_tokens": 322, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + 322, + id="reasoning-already-inside-completion", + ), + ], +) +def test_azure_ai_counts_reasoning_tokens_once_whatever_the_deployment_is_named( + respx_mock: respx.MockRouter, + upstream_model: str, + usage: dict[str, object], + expected_completion_tokens: int, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + respx_mock.post(url__startswith=api_url).mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-x", + "object": "chat.completion", + "created": 1700000000, + "model": upstream_model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "42"}, "finish_reason": "stop"}], + "usage": usage, + }, + ) + ) + + response: Final = litellm.completion( + model="azure_ai/team-reasoner", + messages=[{"role": "user", "content": "What is the meaning of life?"}], + api_base="https://example-resource.services.ai.azure.com/models", + api_key="azure-ai-test-key", + ) + + assert response.usage.completion_tokens == expected_completion_tokens + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + @pytest.mark.respx(assert_all_called=True) def test_azure_ai_grok_stream_folds_reasoning_tokens_into_the_usage_chunk(respx_mock: respx.MockRouter) -> None: api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" @@ -419,6 +479,61 @@ def test_azure_ai_grok_stream_folds_reasoning_tokens_into_the_usage_chunk(respx_ assert usage_chunks[0].total_tokens == usage_chunks[0].prompt_tokens + usage_chunks[0].completion_tokens +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("completion_tokens", "expected_completion_tokens"), + [ + pytest.param(10, 322, id="reasoning-outside-completion"), + pytest.param(322, 322, id="reasoning-already-inside-completion"), + ], +) +def test_azure_ai_stream_counts_reasoning_tokens_once_whatever_the_chunks_call_the_model( + respx_mock: respx.MockRouter, + completion_tokens: int, + expected_completion_tokens: int, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + chunk_base: Final = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "team-reasoner", + } + chunks_sent: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "42"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + "choices": [], + "usage": { + "prompt_tokens": 14, + "completion_tokens": completion_tokens, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + }, + ) + stream_body: Final = "".join(f"data: {json.dumps({**chunk_base, **chunk})}\n\n" for chunk in chunks_sent) + respx_mock.post(url__startswith=api_url).mock( + return_value=Response(200, text=f"{stream_body}data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + ) + + chunks: Final = tuple( + litellm.completion( + model="azure_ai/team-reasoner", + messages=[{"role": "user", "content": "What is the meaning of life?"}], + api_base="https://example-resource.services.ai.azure.com/models", + api_key="azure-ai-test-key", + stream=True, + stream_options={"include_usage": True}, + ) + ) + usage_chunks: Final = [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + assert len(usage_chunks) == 1 + assert usage_chunks[0].completion_tokens == expected_completion_tokens + assert usage_chunks[0].total_tokens == usage_chunks[0].prompt_tokens + usage_chunks[0].completion_tokens + + def test_azure_model_router_response_shows_actual_model(): """ Test that Azure Model Router returns the actual model used in the response, diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 866bdf173a9..176dcfc86c9 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -7,13 +7,17 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +from typing import Final from unittest.mock import patch +import httpx import pytest +import litellm from litellm.llms.azure_ai.anthropic.messages_transformation import ( AzureAnthropicMessagesConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -470,3 +474,80 @@ def test_azure_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_fl and info.get("supports_mid_conversation_system") is not True ] assert missing == [] + + +AUTO_MODE_SAFEGUARDS: Final = ({"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}},) + + +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("client_headers", "optional_params", "expected_beta"), + [ + ({}, {"safeguards": AUTO_MODE_SAFEGUARDS}, "dangerous-tool-use-2026-09-03"), + ( + {"anthropic-beta": "context-1m-2025-08-07"}, + {"safeguards": AUTO_MODE_SAFEGUARDS}, + "context-1m-2025-08-07,dangerous-tool-use-2026-09-03", + ), + ({"anthropic-beta": "context-1m-2025-08-07"}, {}, "context-1m-2025-08-07"), + ({}, {}, None), + ], +) +def test_safeguards_request_carries_the_dangerous_tool_use_beta_without_the_client_header( + client_headers: dict[str, str], optional_params: dict[str, object], expected_beta: str | None +) -> None: + headers, _ = AzureAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=client_headers, + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Run ls -la with the Bash tool"}], + optional_params=optional_params, + litellm_params={"api_key": "test-api-key"}, + ) + assert headers.get("anthropic-beta") == expected_beta, headers + + +def _foundry_answer(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + json={ + "id": "msg_safeguards", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello from the gateway"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_outgoing_azure_request_sends_safeguards_with_the_dangerous_tool_use_beta() -> None: + received: Final[list[httpx.Request]] = [] + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _foundry_answer(request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=[{"role": "user", "content": "Use the Bash tool to run: echo hello from the gateway"}], + model="azure_ai/claude-sonnet-4-6", + api_key="test-api-key", + api_base="https://test-resource.services.ai.azure.com", + safeguards=AUTO_MODE_SAFEGUARDS, + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + sent: Final = received[0] + assert "dangerous-tool-use-2026-09-03" in sent.headers["anthropic-beta"].split(","), dict(sent.headers) + assert json.loads(sent.content)["safeguards"] == list(AUTO_MODE_SAFEGUARDS) diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index d0219ea2963..924b192a00e 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -1,9 +1,12 @@ +import json from collections.abc import Mapping from typing import Final +import httpx import pytest from pydantic import TypeAdapter, ValidationError +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( ir_to_systemone_request, ir_to_systemone_response, @@ -15,10 +18,13 @@ from litellm.types.decisions import ( MAX_DECISION_QUESTIONS, DecisionsIRRefusal, DecisionsIRRequest, + DecisionsIRScoreAnswer, DecisionsRequestBody, OpenAIDecisionRequestBody, UnsupportedDecisionsRequest, ) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager _SYSTEMONE_BODY: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) _OPENAI_BODY: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody) @@ -264,3 +270,66 @@ def test_a_box_answer_without_confidence_reaches_the_caller_unchanged() -> None: response: Final = ir_to_systemone_response(parse_systemone_response(payload, ir), ir) assert response.model_dump(mode="json") == payload + + +@pytest.mark.parametrize("provider", ("nace", "openrouter", "fireworks_ai", "typesafe")) +@pytest.mark.parametrize("legend", (["low", "high"], {"0": "low", "1": "high"})) +def test_score_legends_preserve_labels_probabilities_and_metadata(provider: str, legend: object) -> None: + config: Final = ProviderConfigManager.get_provider_decisions_config("drex-v2.0", LlmProviders(provider)) + assert config is not None + request: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate( + { + "state": {"text": "The order arrived damaged"}, + "questions": {"severity": {"type": "score", "criteria": ["low", "high"]}}, + } + ) + ) + payload: Final = { + "model": "drex-v2.0", + "answers": { + "severity": { + "type": "score", + "score": 0.75, + "confidence": 0.8, + "legend": legend, + "probabilities": {"0": 0.25, "1": 0.75}, + } + }, + "usage": {"input_tokens": 20, "output_tokens": 30}, + "request_id": "test-request", + "evaluation_time_ms": 8.5, + } + original: Final = json.dumps(payload) + response: Final = config.parse_response(payload, request) + answer: Final = response.answers[0] + assert isinstance(answer, DecisionsIRScoreAnswer) + assert (answer.score, answer.confidence) == (0.75, 0.8) + assert tuple((item.value, item.label, item.probability) for item in answer.probabilities) == ( + (0, "low", 0.25), + (1, "high", 0.75), + ) + assert response.model == payload["model"] + assert (response.usage.input_tokens, response.usage.output_tokens) == (20, 30) + assert response.extra == {"request_id": "test-request", "evaluation_time_ms": 8.5} + assert json.dumps(payload) == original + + +@pytest.mark.parametrize( + "payload", ([], {"answers": []}, {"answers": {"a": None}}, {"answers": {"a": {"type": "score", "legend": [1]}}}) +) +@pytest.mark.parametrize("provider", ("nace", "openrouter", "fireworks_ai", "typesafe")) +def test_invalid_provider_response_is_an_api_error(payload: object, provider: str) -> None: + config: Final = ProviderConfigManager.get_provider_decisions_config("drex-v2.0", LlmProviders(provider)) + assert config is not None + request: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate( + { + "state": "text", + "questions": {"a": {"type": "score", "criteria": ["low", "high"]}}, + } + ) + ) + with pytest.raises(BaseLLMException, match="unexpected response") as error: + config.transform_decisions_response("drex-v2.0", provider, httpx.Response(200, json=payload), request) + assert error.value.status_code == 500 diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py index f0f0f9160fb..de4b135e6d0 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py @@ -115,3 +115,22 @@ def test_tool_config_injection_point_is_neither_placed_nor_credited(local_model_ assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["f"] assert "cachePoint" not in json.dumps(request) assert GATEWAY_INJECTED_CACHE_METADATA_KEY not in metadata + + +@pytest.mark.parametrize("value", ["b c", "bad!"], ids=["nbsp", "outside-bedrock-charset"]) +def test_caller_request_metadata_is_ignored_because_invoke_model_has_no_field_for_it(local_model_cost_map, value): + caller_metadata = {"ticket": value} + optional_params = AmazonInvokeNovaConfig().map_openai_params( + non_default_params={"max_tokens": 8, "requestMetadata": caller_metadata}, + optional_params={}, + model=MODEL, + drop_params=False, + ) + + request = _transform_request( + messages=[{"role": "user", "content": "hi"}], + optional_params={**optional_params, "requestMetadata": caller_metadata}, + ) + + assert "requestMetadata" not in request + assert request["inferenceConfig"] == {"maxTokens": 8} diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 16f44f4f35f..ca34a4eee06 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -3,11 +3,13 @@ import copy import json import uuid from types import SimpleNamespace +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. @@ -482,7 +484,7 @@ def test_output_config_format_converted_for_bedrock_chat_invoke_request(): with patch( # test-quality-ok: pin non-native path "litellm.llms.bedrock.common_utils._bedrock_model_supports", - side_effect=lambda _model, key: key == "supports_output_config", + side_effect=lambda _model, key, *_: key == "supports_output_config", ): result = config.transform_request( model="anthropic.claude-opus-4-7", @@ -504,10 +506,12 @@ def test_output_config_format_converted_for_bedrock_chat_invoke_request(): @pytest.mark.parametrize("model", ["anthropic.claude-opus-4-7", "us.anthropic.claude-opus-4-8"]) -def test_output_config_format_inlined_for_bedrock_chat_invoke_opus_4_7_and_4_8(local_model_cost_map, model): - """Bedrock rejects ``output_config.format`` on Claude Opus 4.7 and 4.8, so the - Invoke chat path inlines the schema into the last user message and keeps effort, - driven by the cost map alone (no capability stub).""" +def test_output_config_format_inlined_for_bedrock_chat_invoke_profile_without_native_output( + set_native_structured_output, model +): + """On a profile without native structured output the Invoke chat path inlines the + schema into the last user message and keeps effort.""" + set_native_structured_output(model, None) schema = {"type": "object", "properties": {"answer": {"type": "string"}}} result = AmazonAnthropicClaudeConfig().transform_request( @@ -525,12 +529,155 @@ def test_output_config_format_inlined_for_bedrock_chat_invoke_opus_4_7_and_4_8(l assert json.loads(result["messages"][-1]["content"][-1]["text"]) == schema +TREE_SCHEMA: Final = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": {"children": {"type": "array", "items": {"$ref": "#/$defs/Node"}}}, + } + }, +} +OPEN_OBJECT_SCHEMA: Final = { + "type": "object", + "properties": {"counts": {"type": "object", "additionalProperties": {"type": "integer"}}}, +} + + +def _invoke_params_for_schema( + model: str, schema: Mapping[str, object], thinking_params: Mapping[str, object] | None = None +) -> Mapping[str, object]: + return AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": {"type": "json_schema", "json_schema": {"name": "out", "schema": schema}}, + **({"thinking": {"type": "adaptive"}} if thinking_params is None else thinking_params), + }, + optional_params={}, + model=model, + drop_params=False, + ) + + +@pytest.mark.parametrize( + ("native_key", "thinking_params", "native"), + [ + ("legacy", {}, True), + ("bedrock", {}, False), + ("bedrock", {"thinking": {"type": "adaptive"}}, True), + ("bedrock", {"reasoning_effort": "low"}, True), + ], + ids=["legacy-key", "bedrock-key-no-thinking", "bedrock-key-adaptive", "bedrock-key-effort"], +) +def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( + set_native_structured_output, native_key, thinking_params, native +): + model = "us.anthropic.claude-haiku-5-5" + set_native_structured_output(model, native_key) + schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]} + + result = _invoke_params_for_schema(model, schema, thinking_params) + + assert ("output_format" in result) is native + assert ("tools" in result) is not native + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +@pytest.mark.parametrize("schema", [TREE_SCHEMA, OPEN_OBJECT_SCHEMA], ids=["recursive", "open-object"]) +def test_schema_native_output_cannot_hold_keeps_the_tool_path_on_a_native_profile( + set_native_structured_output, schema, native_key +): + model = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + result = _invoke_params_for_schema(model, schema) + + assert "output_format" not in result + assert [tool["name"] for tool in result["tools"]] == ["json_tool_call"] + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_closed_schema_goes_native_on_either_native_key(set_native_structured_output, native_key): + model = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + schema = {"type": "object", "properties": {"age": {"type": "integer", "minimum": 0}}, "required": ["age"]} + + result = _invoke_params_for_schema(model, schema) + + assert result["output_format"]["schema"]["properties"]["age"]["type"] == "integer" + assert "minimum" not in result["output_format"]["schema"]["properties"]["age"] + assert "tools" not in result + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_shared_defs_reach_native_output_once_instead_of_inlined_per_use(set_native_structured_output, native_key): + model = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + party_ref = {"$ref": "#/$defs/Party"} + schema = { + "type": "object", + "properties": {"buyer": party_ref, "seller": party_ref}, + "required": ["buyer", "seller"], + "$defs": { + "Party": { + "type": "object", + "properties": {"email": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": ["email"], + } + }, + } + + result = _invoke_params_for_schema(model, schema) + + sent = result["output_format"]["schema"] + assert sent["properties"]["buyer"] == party_ref + assert json.dumps(sent).count('"anyOf"') == 1 + assert "tools" not in result + + + +class Relabel(BaseModel): + model_config = ConfigDict(frozen=True) + + box: str + new_label: str + + +class PuzzleAnswer(BaseModel): + model_config = ConfigDict(frozen=True) + + relabels: list[Relabel] + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_nested_pydantic_response_format_reaches_native_output_with_refs_that_resolve( + set_native_structured_output, native_key +): + model = "us.anthropic.claude-haiku-5-5" + set_native_structured_output(model, native_key) + config: Final = AmazonAnthropicClaudeConfig() + + result: Final = config.map_openai_params( + non_default_params={ + "response_format": config.get_json_schema_from_pydantic_object(PuzzleAnswer), + "thinking": {"type": "adaptive"}, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + sent: Final = result["output_format"]["schema"] + item_ref: Final = sent["properties"]["relabels"]["items"]["$ref"] + assert item_ref.removeprefix("#/$defs/") in sent["$defs"] + assert "tools" not in result + def test_output_config_format_forwarded_for_bedrock_chat_invoke_request(local_model_cost_map): """Bedrock Invoke chat path forwards ``output_config.format`` alongside effort for models with native structured-output support (Claude Sonnet 4.6).""" schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}, "additionalProperties": False}, } result = AmazonAnthropicClaudeConfig().transform_request( @@ -671,7 +818,7 @@ def test_bedrock_chat_invoke_drop_params_keeps_native_output_config_format(local monkeypatch.setattr(litellm, "drop_params", True) schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}, "additionalProperties": False}, } result = AmazonAnthropicClaudeConfig().transform_request( @@ -713,10 +860,11 @@ def test_bedrock_chat_invoke_drop_params_still_inlines_for_non_native(local_mode "model", ["us.anthropic.claude-fable-5-1", "anthropic.claude-fable-5-1"], ) -def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice(local_model_cost_map, model): - """Regression: Bedrock rejects both native ``output_config.format`` and forced - tool_choice for Fable 5.1, so invoke must use the tool-based path without a +def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice(set_native_structured_output, model): + """Regression: Bedrock rejects forced tool_choice for Fable 5.1, so on a profile + without native structured output invoke must use the tool-based path without a forced ``tool_choice``.""" + set_native_structured_output(model, None) result = AmazonAnthropicClaudeConfig().map_openai_params( non_default_params={ "response_format": { @@ -738,7 +886,10 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice @pytest.mark.parametrize("model", ["us.anthropic.claude-sonnet-5", "us.anthropic.claude-fable-5-1"]) -def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map, model): +def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking( + set_native_structured_output, model +): + set_native_structured_output(model, None) result = AmazonAnthropicClaudeConfig().map_openai_params( non_default_params={ "response_format": { @@ -761,9 +912,10 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th assert result["output_config"] == {"effort": "high"} -def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_tool(local_model_cost_map): +def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_tool(set_native_structured_output): """Anthropic accepts a forced tool_choice next to adaptive thinking, and dropping it lets the model answer in prose instead of the requested JSON.""" + set_native_structured_output("us.anthropic.claude-opus-5", None) result = AmazonAnthropicClaudeConfig().map_openai_params( non_default_params={ "response_format": { diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 6c996a7ce75..6092934fc7a 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -3,7 +3,7 @@ import copy import json import os import re -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final from unittest.mock import MagicMock, Mock, patch @@ -12,12 +12,18 @@ import pytest import respx import litellm -from litellm import ModelResponse, completion +from litellm import ModelResponse, Router, completion from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, _bedrock_tools_pt, ) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import CONVERTED_SYSTEM_NOTE +from litellm.llms.anthropic.chat.transformation import ( + NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES, + NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS, + NATIVE_OUTPUT_MAX_UNION_PARAMETERS, +) +from litellm.llms.bedrock.common_utils import BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -1246,7 +1252,9 @@ async def test_transform_request_replays_redacted_thinking_only_to_claude( """A model id that names its model decides on its own. An application inference profile ARN hides its model, so the deployment's ``base_model`` is what marks it as Claude, and without one the redacted block is dropped the way it is for every non-Claude target.""" - redacted_data: Final = base64.b64encode(b"EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP").decode() + redacted_data: Final = base64.b64encode( + b"EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP" + ).decode() messages: Final = [ {"role": "user", "content": "hi"}, { @@ -1475,32 +1483,25 @@ def test_output_config_format_dropped_on_unsupported_converse_model_warns(caplog """When Converse model lacks native structured-output support, the silently dropped ``output_config.format`` must surface as a warning so callers can diagnose plain-text responses.""" - from unittest.mock import patch - config = AmazonConverseConfig() schema = { "type": "object", "properties": {"answer": {"type": "string"}}, } - with patch.object( - AmazonConverseConfig, - "_supports_native_structured_outputs", - return_value=False, - ): - with caplog.at_level("WARNING"): - result = config._transform_request( - model="bedrock/converse/us.anthropic.claude-3-haiku-20240307-v1:0", - messages=[{"role": "user", "content": "hi"}], - optional_params={ - "maxTokens": 256, - "output_config": { - "format": {"type": "json_schema", "schema": schema}, - }, + with caplog.at_level("WARNING"): + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "output_config": { + "format": {"type": "json_schema", "schema": schema}, }, - litellm_params={}, - headers={}, - ) + }, + litellm_params={}, + headers={}, + ) assert "outputConfig" not in result assert any("dropping `output_config.format`" in record.getMessage() for record in caplog.records) @@ -4198,7 +4199,7 @@ def test_request_metadata_key_constraints(): # Test empty key invalid_metadata = {"": "value"} - with pytest.raises(Exception, match=r"(?i)key length|empty"): + with pytest.raises(Exception, match=r"key '' must be 1-256 characters"): config._transform_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", messages=messages, @@ -4240,6 +4241,18 @@ def test_request_metadata_value_constraints(): ) +@pytest.mark.parametrize("metadata", [{"ticket": "a\u00a0b"}, {"a\u2003b": "ticket"}]) +def test_request_metadata_rejects_whitespace_outside_ascii(metadata): + with pytest.raises(litellm.exceptions.BadRequestError, match=r"must be [01]-256 characters of"): + AmazonConverseConfig()._transform_request( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + optional_params={"requestMetadata": metadata}, + litellm_params={}, + headers={}, + ) + + def test_request_metadata_character_pattern(): """Test character pattern validation for keys and values.""" config = AmazonConverseConfig() @@ -4841,24 +4854,6 @@ def test_translate_response_format_native_output_config(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) -BEDROCK_OPUS_4_7_AND_4_8_MODELS: Final = ( - "anthropic.claude-opus-4-7", - "global.anthropic.claude-opus-4-7", - "us.anthropic.claude-opus-4-7", - "eu.anthropic.claude-opus-4-7", - "au.anthropic.claude-opus-4-7", - "jp.anthropic.claude-opus-4-7", - "anthropic.claude-opus-4-8", - "global.anthropic.claude-opus-4-8", - "us.anthropic.claude-opus-4-8", - "eu.anthropic.claude-opus-4-8", - "au.anthropic.claude-opus-4-8", - "jp.anthropic.claude-opus-4-8", - "us-gov.anthropic.claude-opus-4-8", - "us-gov-west-1/anthropic.claude-opus-4-8", - "us-gov-east-1/anthropic.claude-opus-4-8", -) - CAPITAL_RESPONSE_FORMAT: Final = { "type": "json_schema", "json_schema": { @@ -4873,10 +4868,15 @@ CAPITAL_RESPONSE_FORMAT: Final = { } -def _converse_request_for_json_schema(model: str, stream: bool) -> tuple[dict, dict]: +ADAPTIVE_THINKING: Final = {"type": "adaptive"} + + +def _converse_request_for_json_schema( + model: str, stream: bool, thinking_params: Mapping[str, object] | None = None +) -> tuple[Mapping[str, object], RequestObject]: config = AmazonConverseConfig() optional_params = config.map_openai_params( - non_default_params={"response_format": CAPITAL_RESPONSE_FORMAT, "stream": stream}, + non_default_params={"response_format": CAPITAL_RESPONSE_FORMAT, "stream": stream, **(thinking_params or {})}, optional_params={}, model=model, drop_params=False, @@ -4891,21 +4891,452 @@ def _converse_request_for_json_schema(model: str, stream: bool) -> tuple[dict, d return optional_params, request -@pytest.mark.parametrize("model", BEDROCK_OPUS_4_7_AND_4_8_MODELS) +@pytest.mark.parametrize( + ("native_profile", "tool_profile"), + [ + ("us.anthropic.claude-opus-4-8", "anthropic.claude-opus-4-8"), + ("global.anthropic.claude-opus-4-7", "us.anthropic.claude-opus-4-7"), + ("anthropic.claude-opus-4-8", "us-gov.anthropic.claude-opus-4-8"), + ], +) @pytest.mark.parametrize("stream", [False, True]) -def test_opus_4_7_and_4_8_json_schema_sent_as_forced_tool_not_output_config(monkeypatch, model, stream): - """Regression for issue #27846: Bedrock rejects outputConfig on Opus 4.7 and 4.8 - (``output_config.format: Extra inputs are not permitted``), so json_schema has to - go out as the forced json_tool_call tool, streamed through fake_stream.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +def test_json_schema_shape_follows_each_inference_profile_own_native_output_flag( + set_native_structured_output, native_profile, tool_profile, stream +): + set_native_structured_output(native_profile, "bedrock") + set_native_structured_output(tool_profile, None) - optional_params, request = _converse_request_for_json_schema(model=model, stream=stream) + thinking_params: Final = {"thinking": ADAPTIVE_THINKING} + native_params, native_request = _converse_request_for_json_schema(native_profile, stream, thinking_params) + tool_params, tool_request = _converse_request_for_json_schema(tool_profile, stream, thinking_params) + + assert native_request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "capital" + assert "toolConfig" not in native_request + assert "fake_stream" not in native_params + assert "outputConfig" not in tool_request + assert [tool["toolSpec"]["name"] for tool in tool_request["toolConfig"]["tools"]] == ["json_tool_call"] + assert tool_request["toolConfig"]["toolChoice"] == {"tool": {"name": "json_tool_call"}} + assert tool_params.get("fake_stream", False) is stream + + +@pytest.mark.parametrize( + ("native_key", "thinking_params", "native"), + [ + ("legacy", {}, True), + ("bedrock", {}, False), + ("bedrock", {"thinking": ADAPTIVE_THINKING}, True), + ("bedrock", {"thinking": {"type": "enabled", "budget_tokens": 1024}}, True), + ("bedrock", {"reasoning_effort": "low"}, True), + ("bedrock", {"reasoning_effort": "none"}, False), + ("bedrock", {"thinking": {"type": "disabled"}}, False), + ], + ids=[ + "legacy-key", + "bedrock-key-no-thinking", + "bedrock-key-adaptive", + "bedrock-key-enabled", + "bedrock-key-effort", + "bedrock-key-effort-none", + "bedrock-key-thinking-disabled", + ], +) +def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( + set_native_structured_output, native_key, thinking_params, native +): + model: Final = "us.anthropic.claude-haiku-5-5" + set_native_structured_output(model, native_key) + + _, request = _converse_request_for_json_schema(model, False, thinking_params) + + assert ("outputConfig" in request) is native + assert ("toolConfig" in request) is not native + + +def test_adaptive_thinking_is_kept_next_to_json_schema_on_a_native_output_profile(set_native_structured_output): + model: Final = "us.anthropic.claude-opus-5" + set_native_structured_output(model, "bedrock") + config: Final = AmazonConverseConfig() + thinking: Final = {"type": "adaptive", "display": "summarized"} + + optional_params: Final = config.map_openai_params( + non_default_params={"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": thinking}, + optional_params={}, + model=model, + drop_params=False, + ) + request: Final = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Name the capital of France."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["additionalModelRequestFields"]["thinking"] == thinking + assert request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "capital" + assert "toolConfig" not in request + + +_OPUS_4_8_INVOKE_RESPONSE: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": '{"city": "Paris", "country": "France"}'}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +def _routed_request_body(router: Router, model_name: str, route: str) -> Mapping[str, object]: + response: Final = _mock_converse_response() + if route == "invoke": + response.json.return_value = _OPUS_4_8_INVOKE_RESPONSE + response.text = json.dumps(_OPUS_4_8_INVOKE_RESPONSE) + client: Final = HTTPHandler() + with patch.object(client, "post", return_value=response) as mock_post: + router.completion( + model=model_name, + messages=[{"role": "user", "content": "Name the capital of France."}], + response_format=CAPITAL_RESPONSE_FORMAT, + thinking=ADAPTIVE_THINKING, + client=client, + ) + return json.loads(mock_post.call_args.kwargs["data"]) + + +def _sends_native_output(body: Mapping[str, object], route: str) -> bool: + if route == "converse": + return "outputConfig" in body and "toolConfig" not in body + output_config: Final = body.get("output_config") + return isinstance(output_config, dict) and "format" in output_config and "tools" not in body + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_a_deployment_native_output_opt_in_stays_on_that_deployment(set_native_structured_output, route): + set_native_structured_output("us.anthropic.claude-opus-4-8", None) + litellm_params: Final = { + "model": f"bedrock/{route}/us.anthropic.claude-opus-4-8", + "aws_access_key_id": "fake-access-key", + "aws_secret_access_key": "fake-secret-key", + "aws_region_name": "us-east-1", + } + router: Final = Router( + model_list=[ + { + "model_name": "opted-in", + "litellm_params": litellm_params, + "model_info": {"supports_bedrock_native_structured_output": True}, + }, + {"model_name": "sibling", "litellm_params": litellm_params}, + ] + ) + + assert _sends_native_output(_routed_request_body(router, "opted-in", route), route) + assert not _sends_native_output(_routed_request_body(router, "sibling", route), route) + + +NATIVE_OUTPUT_REJECTED_KEYWORDS: Final = frozenset({"minimum", "maximum", "uniqueItems", "oneOf"}) +CONSTRAINED_PERSON_SCHEMA: Final = { + "type": "object", + "properties": { + "age": {"type": "integer", "minimum": 0, "maximum": 150}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 2, "uniqueItems": True}, + "aliases": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "contact": {"oneOf": [{"type": "string"}, {"type": "integer"}]}, + "address": {"type": ["object", "null"], "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + "required": ["age", "tags", "aliases", "contact", "address"], +} +RECURSIVE_TREE_SCHEMA: Final = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "required": ["root"], + "$defs": { + "Node": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "children": {"type": "array", "items": {"$ref": "#/$defs/Node"}}, + }, + "required": ["name", "children"], + } + }, +} +OPEN_OBJECT_SCHEMA: Final = { + "type": "object", + "properties": {"counts": {"type": "object", "additionalProperties": {"type": "integer"}}}, + "required": ["counts"], +} +LOOKAROUND_PATTERN_SCHEMA: Final = { + "type": "object", + "properties": {"amount": {"type": "string", "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$"}}, + "required": ["amount"], +} +UNKNOWN_FORMAT_SCHEMA: Final = { + "type": "object", + "properties": {"blob": {"type": "string", "format": "x-litellm-unlisted-format"}}, + "required": ["blob"], +} +ROOT_REF_SCHEMA: Final = { + "type": "object", + "properties": {"name": {"type": "string"}, "parent": {"anyOf": [{"$ref": "#"}, {"type": "null"}]}}, + "required": ["name", "parent"], +} + + +def _schema_keys(node: object) -> frozenset[str]: + if isinstance(node, dict): + return frozenset(node).union(*(_schema_keys(value) for value in node.values())) + if isinstance(node, list): + return frozenset().union(*(_schema_keys(item) for item in node)) + return frozenset() + + +def _converse_request_for_schema(model: str, schema: Mapping[str, object]) -> RequestObject: + config: Final = AmazonConverseConfig() + optional_params: Final = config.map_openai_params( + non_default_params={ + "response_format": {"type": "json_schema", "json_schema": {"name": "out", "schema": schema}}, + "thinking": ADAPTIVE_THINKING, + }, + optional_params={}, + model=model, + drop_params=False, + ) + return config._transform_request( + model=model, + messages=[{"role": "user", "content": "Fill in a plausible example."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_native_output_schema_drops_keywords_bedrock_rejects_and_closes_objects( + set_native_structured_output, native_key +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, CONSTRAINED_PERSON_SCHEMA) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert not _schema_keys(sent) & NATIVE_OUTPUT_REJECTED_KEYWORDS + assert sent["additionalProperties"] is False + assert sent["properties"]["contact"]["anyOf"] == [{"type": "string"}, {"type": "integer"}] + assert "minimum value: 0" in sent["properties"]["age"]["description"] + assert "minItems" not in sent["properties"]["tags"] + assert sent["properties"]["aliases"]["minItems"] == 1 + assert sent["properties"]["address"]["additionalProperties"] is False + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +@pytest.mark.parametrize( + ("schema", "kept_keyword"), + [ + (RECURSIVE_TREE_SCHEMA, "$ref"), + (ROOT_REF_SCHEMA, "$ref"), + (OPEN_OBJECT_SCHEMA, "additionalProperties"), + (LOOKAROUND_PATTERN_SCHEMA, "pattern"), + (UNKNOWN_FORMAT_SCHEMA, "format"), + ], + ids=["recursive", "root-ref", "open-object", "lookaround-pattern", "unlisted-format"], +) +def test_schema_native_output_cannot_hold_keeps_the_tool_path_on_a_native_profile( + set_native_structured_output, schema, kept_keyword, native_key +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, schema) assert "outputConfig" not in request - assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["json_tool_call"] - assert request["toolConfig"]["toolChoice"] == {"tool": {"name": "json_tool_call"}} - assert optional_params.get("fake_stream", False) is stream + tool_spec: Final = request["toolConfig"]["tools"][0]["toolSpec"] + assert tool_spec["name"] == "json_tool_call" + assert kept_keyword in _schema_keys(tool_spec["inputSchema"]["json"]) + + +def _shared_defs_schema(defs_key: str) -> Mapping[str, object]: + party_ref: Final = {"$ref": f"#/{defs_key}/Party"} + return { + "type": "object", + "properties": {"buyer": party_ref, "seller": party_ref, "broker": {"anyOf": [party_ref, {"type": "null"}]}}, + "required": ["buyer", "seller", "broker"], + defs_key: { + "Party": { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}, "email": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": ["name", "email"], + } + }, + } + + +@pytest.mark.parametrize("defs_key", ["$defs", "definitions"]) +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_shared_defs_reach_native_output_once_instead_of_inlined_per_use( + set_native_structured_output, native_key, defs_key +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, _shared_defs_schema(defs_key)) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert sent["properties"]["buyer"] == {"$ref": f"#/{defs_key}/Party"} + assert json.dumps(sent).count('"anyOf"') == 2 + assert sent[defs_key]["Party"]["additionalProperties"] is False + + +def _flat_schema(property_schema: Mapping[str, object], count: int, required: bool) -> Mapping[str, object]: + names: Final = tuple(f"field_{index}" for index in range(count)) + return { + "type": "object", + "properties": dict.fromkeys(names, property_schema), + "required": list(names) if required else [], + } + + +def _optional_groups_schema(count: int, group_size: int = 6) -> Mapping[str, object]: + groups: Final = { + f"group_{start}": _flat_schema({"type": "string"}, min(group_size, count - start), required=False) + for start in range(0, count, group_size) + } + return {"type": "object", "properties": groups, "required": list(groups)} + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +@pytest.mark.parametrize( + ("schema_for", "limit"), + [ + (lambda count: _flat_schema({"type": ["string", "null"]}, count, True), NATIVE_OUTPUT_MAX_UNION_PARAMETERS), + ( + lambda count: _flat_schema({"anyOf": [{"type": "string"}, {"type": "integer"}]}, count, True), + NATIVE_OUTPUT_MAX_UNION_PARAMETERS, + ), + (_optional_groups_schema, NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS), + ( + lambda count: _flat_schema({"type": "string"}, count, False), + NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES.bit_length() - 1, + ), + ], + ids=["type-list-unions", "any-of-unions", "optional", "optional-grammar-branches"], +) +def test_schema_past_a_native_output_size_limit_keeps_the_tool_path( + set_native_structured_output, native_key, schema_for, limit +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + at_limit: Final = _converse_request_for_schema(model, schema_for(limit)) + past_limit: Final = _converse_request_for_schema(model, schema_for(limit + 1)) + + sent: Final = json.loads(at_limit["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert sent["properties"].keys() == schema_for(limit)["properties"].keys() + assert "outputConfig" not in past_limit + assert past_limit["toolConfig"]["tools"][0]["toolSpec"]["name"] == "json_tool_call" + + +@pytest.mark.parametrize( + ("native_key", "native_past_limit"), [("legacy", True), ("bedrock", False)], ids=["legacy-key", "bedrock-key"] +) +def test_field_count_limit_sends_the_json_tool_only_on_the_bedrock_native_key( + set_native_structured_output, native_key, native_past_limit +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + at_limit: Final = _converse_request_for_schema( + model, _flat_schema({"type": "string"}, BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES, True) + ) + past_limit: Final = _converse_request_for_schema( + model, _flat_schema({"type": "string"}, BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES + 1, True) + ) + + assert "outputConfig" in at_limit + assert ("outputConfig" in past_limit) is native_past_limit + + +NUMERIC_FORMAT_SCHEMA: Final = { + "type": "object", + "properties": { + "count": {"type": "integer", "format": "int32"}, + "total": {"type": "integer", "format": "int64"}, + "ratio": {"type": "number", "format": "double"}, + }, + "required": ["count", "total", "ratio"], +} + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_numeric_formats_keep_native_output(set_native_structured_output, native_key): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, NUMERIC_FORMAT_SCHEMA) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert {name: field["format"] for name, field in sent["properties"].items()} == { + "count": "int32", + "total": "int64", + "ratio": "double", + } + + +def _doubling_defs_schema(depth: int) -> Mapping[str, object]: + levels: Final = { + f"Level{index}": { + "type": "object", + "properties": { + "left": {"$ref": f"#/$defs/Level{index + 1}"}, + "right": {"$ref": f"#/$defs/Level{index + 1}"}, + }, + "required": ["left", "right"], + } + for index in range(depth) + } + return { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Level0"}}, + "required": ["root"], + "$defs": {**levels, f"Level{depth}": {"type": "string"}}, + } + + +@pytest.mark.timeout(30) +def test_defs_that_double_at_every_level_reach_native_output_without_being_expanded(set_native_structured_output): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, "legacy") + schema: Final = _doubling_defs_schema(30) + + request: Final = _converse_request_for_schema(model, schema) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert sent["$defs"].keys() == schema["$defs"].keys() + assert sent["properties"]["root"] == {"$ref": "#/$defs/Level0"} + + +@pytest.mark.parametrize( + ("native_key", "forwarded"), [("legacy", True), ("bedrock", False)], ids=["legacy-key", "bedrock-key"] +) +def test_recursive_output_config_format_follows_the_profile_native_key( + set_native_structured_output, native_key, forwarded +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = AmazonConverseConfig()._transform_request( + model=model, + messages=[{"role": "user", "content": "Fill in a plausible example."}], + optional_params={"output_config": {"format": {"type": "json_schema", "schema": RECURSIVE_TREE_SCHEMA}}}, + litellm_params={}, + headers={}, + ) + + assert ("outputConfig" in request) is forwarded + def test_sonnet_4_6_json_schema_still_uses_native_output_config(monkeypatch): @@ -8008,10 +8439,13 @@ def test_unforced_tool_choice_unaffected_on_fable_5_1_converse(local_model_cost_ "model", ["anthropic.claude-fable-5-1", "us.anthropic.claude-fable-5-1"], ) -def test_response_format_avoids_native_and_forced_tool_choice_on_fable_5_1_converse(local_model_cost_map, model): - """Regression: Bedrock rejects both ``outputConfig`` structured output and forced - tool_choice for Fable 5.1, so response_format must map to a tool without a forced - tool_choice.""" +def test_response_format_avoids_native_and_forced_tool_choice_on_fable_5_1_converse( + set_native_structured_output, model +): + """Regression: Bedrock rejects forced tool_choice for Fable 5.1, so on a profile + without native structured output response_format must map to a tool without a + forced tool_choice.""" + set_native_structured_output(model, None) config = AmazonConverseConfig() result = config.map_openai_params( diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 425beefac8a..b050bb7625b 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1760,6 +1760,23 @@ class TestBedrockBatchNonChatEndpointRecords: assert "metadata" not in model_input assert "requestMetadata" not in model_input + def test_nova_record_request_metadata_stays_out_of_the_bedrock_body(self): + model_input = self._transform( + { + "custom_id": "4d", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": self.NOVA_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "requestMetadata": {"ticket": "b\u00a0c"}, + }, + } + ) + + assert "requestMetadata" not in model_input + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + @pytest.mark.parametrize( "body", [ diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index c91eb75dbd9..eb6b4fdb063 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1050,6 +1050,7 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): "schema": { "type": "object", "properties": {"answer": {"type": "string"}}, + "additionalProperties": False, }, } optional_params = { @@ -1098,7 +1099,7 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): with patch( "litellm.llms.bedrock.common_utils._bedrock_model_supports", - side_effect=lambda _model, key: key == "supports_output_config", + side_effect=lambda _model, key, *_: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1170,7 +1171,7 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema with patch( "litellm.llms.bedrock.common_utils._bedrock_model_supports", - side_effect=lambda _model, key: key == "supports_output_config", + side_effect=lambda _model, key, *_: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1677,7 +1678,12 @@ def test_bedrock_messages_does_not_add_dangerous_tool_use_beta_without_safeguard assert "dangerous-tool-use-2026-09-03" not in result.get("anthropic_beta", []) -_NAME_SCHEMA: Final = {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} +_NAME_SCHEMA: Final = { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, +} @pytest.mark.parametrize( @@ -3216,11 +3222,12 @@ def test_bedrock_messages_forwards_output_config_format_natively(local_model_cos assert "zebra_count" not in json.dumps(result["messages"]) -def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): - """Bedrock rejects ``output_config.format`` for the Claude 5 family, so the - schema falls back to the inline-text path instead of a deterministic 400.""" +def test_bedrock_messages_inlines_schema_for_a_profile_without_native_output(set_native_structured_output): + """On a profile without native structured output the schema falls back to the + inline-text path instead of a deterministic 400.""" from litellm.types.router import GenericLiteLLMParams + set_native_structured_output("us.anthropic.claude-sonnet-5", None) cfg = AmazonAnthropicClaudeMessagesConfig() schema = { "type": "object", @@ -3243,6 +3250,115 @@ def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): assert json.loads(last_content[-1]["text"]) == schema +def _bedrock_messages_request_with_schema(model: str, schema: Mapping[str, object]) -> Mapping[str, object]: + from litellm.types.router import GenericLiteLLMParams + + return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + +def test_bedrock_messages_native_output_drops_schema_keywords_bedrock_rejects(set_native_structured_output): + model: Final = "us.anthropic.claude-sonnet-5" + set_native_structured_output(model, "bedrock") + schema: Final = { + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "ids": {"type": "array", "items": {"type": "string"}, "uniqueItems": True}, + }, + "required": ["count", "ids"], + } + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + sent: Final = result["output_config"]["format"]["schema"] + assert "minimum" not in sent["properties"]["count"] + assert "uniqueItems" not in sent["properties"]["ids"] + assert sent["additionalProperties"] is False + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_bedrock_messages_shared_defs_reach_native_output_once_instead_of_inlined_per_use( + set_native_structured_output, native_key +): + model: Final = "us.anthropic.claude-sonnet-5" + set_native_structured_output(model, native_key) + party_ref: Final = {"$ref": "#/$defs/Party"} + schema: Final = { + "type": "object", + "properties": {"buyer": party_ref, "seller": party_ref}, + "required": ["buyer", "seller"], + "additionalProperties": False, + "$defs": { + "Party": { + "type": "object", + "properties": {"email": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": ["email"], + "additionalProperties": False, + } + }, + } + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + sent: Final = result["output_config"]["format"]["schema"] + assert sent["properties"]["buyer"] == party_ref + assert json.dumps(sent).count('"anyOf"') == 1 + + +MESSAGES_TREE_SCHEMA: Final = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": {"children": {"type": "array", "items": {"$ref": "#/$defs/Node"}}}, + } + }, +} +MESSAGES_OPEN_OBJECT_SCHEMA: Final = { + "type": "object", + "properties": {"counts": {"type": "object", "additionalProperties": {"type": "integer"}}}, +} + + +@pytest.mark.parametrize( + "schema", [MESSAGES_TREE_SCHEMA, MESSAGES_OPEN_OBJECT_SCHEMA], ids=["recursive", "open-object"] +) +def test_bedrock_messages_schema_native_output_cannot_hold_is_inlined_on_a_bedrock_native_profile( + set_native_structured_output, schema +): + model: Final = "us.anthropic.claude-sonnet-5" + set_native_structured_output(model, "bedrock") + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + assert "output_config" not in result + assert json.loads(result["messages"][-1]["content"][-1]["text"]) == schema + + +@pytest.mark.parametrize( + "schema", [MESSAGES_TREE_SCHEMA, MESSAGES_OPEN_OBJECT_SCHEMA], ids=["recursive", "open-object"] +) +def test_bedrock_messages_schema_native_output_cannot_hold_is_forwarded_on_a_legacy_native_profile( + set_native_structured_output, schema +): + model: Final = "us.anthropic.claude-sonnet-4-6" + set_native_structured_output(model, "legacy") + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + assert result["output_config"]["format"]["schema"] == schema + + def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(local_model_cost_map): """When a request carries both schema forms, the legacy top-level ``output_format`` keeps winning, matching the pre-existing precedence.""" @@ -3251,7 +3367,11 @@ def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(lo cfg = AmazonAnthropicClaudeMessagesConfig() legacy_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"legacy_field": {"type": "string"}}}, + "schema": { + "type": "object", + "properties": {"legacy_field": {"type": "string"}}, + "additionalProperties": False, + }, } newer_format = { "type": "json_schema", @@ -3285,7 +3405,7 @@ def test_bedrock_messages_drop_params_keeps_native_output_config_format(local_mo cfg = AmazonAnthropicClaudeMessagesConfig() schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}, "additionalProperties": False}, } result = cfg.transform_anthropic_messages_request( @@ -3310,7 +3430,7 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo cfg = AmazonAnthropicClaudeMessagesConfig() schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}, "additionalProperties": False}, } result = cfg.transform_anthropic_messages_request( diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 57768b41943..85de558a8b2 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -10,17 +10,23 @@ Validates that: """ import copy +import json +from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Literal +from typing import Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock import pytest +import respx +from pydantic import JsonValue +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.noma import NomaV2Guardrail from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -803,11 +809,17 @@ class TestBedrockPassthroughGuardrailHandlerOutput: def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: import json + + return _build_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload, separators=(",", ":")).encode(), + ) + + +def _build_frame(headers: dict[str, str], payload_bytes: bytes) -> bytes: import struct from binascii import crc32 as esm_crc32 - payload_bytes = json.dumps(payload, separators=(",", ":")).encode() - def _encode_str_header(name: str, value: str) -> bytes: name_b = name.encode() value_b = value.encode() @@ -815,11 +827,7 @@ def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b ) - headers_bytes = ( - _encode_str_header(":event-type", event_type) - + _encode_str_header(":content-type", "application/json") - + _encode_str_header(":message-type", "event") - ) + headers_bytes = b"".join(_encode_str_header(name, value) for name, value in headers.items()) headers_length = len(headers_bytes) total_length = 12 + headers_length + len(payload_bytes) + 4 @@ -853,7 +861,10 @@ class _CardMaskingGuardrail(CustomGuardrail): async def _run_post_call_guardrail_on_stream( - stream_bytes: bytes, guardrail: CustomGuardrail, monkeypatch: pytest.MonkeyPatch + stream_bytes: bytes, + guardrail: CustomGuardrail, + monkeypatch: pytest.MonkeyPatch, + endpoint: str = "model/claude/converse-stream", ) -> bytes: monkeypatch.setattr(litellm, "callbacks", [guardrail]) logging_obj = LiteLLMLoggingObj( @@ -872,7 +883,7 @@ async def _run_post_call_guardrail_on_stream( data={ "model": "claude", "custom_llm_provider": "bedrock", - "endpoint": "model/claude/converse-stream", + "endpoint": endpoint, "litellm_logging_obj": logging_obj, }, ) @@ -1342,3 +1353,1334 @@ class TestDeAnonymizeConverseStream: hook_spy.assert_not_called() assert result is stream_bytes + + +def _decoded_frames(raw: bytes) -> list[tuple[str, str, dict]]: + import json + from botocore.eventstream import EventStreamBuffer + + buf = EventStreamBuffer() + buf.add_data(raw) + return [ + ( + str(msg.headers[":message-type"]), + str(msg.headers.get(":event-type") or msg.headers.get(":exception-type")), + json.loads(msg.payload), + ) + for msg in buf + ] + + +_CONVERSE_STREAM = "model/claude/converse-stream" +_INVOKE_STREAM = "model/anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream" +_APPLICATION_PROFILE = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc" + + +async def _chunks(*chunks: bytes): + for chunk in chunks: + yield chunk + + +async def _relay_through_post_call_guardrail( + chunks: tuple[bytes, ...], + guardrail: CustomGuardrail, + monkeypatch: pytest.MonkeyPatch, + holds_until_moderated: bool = False, + received: list[bytes] | None = None, + outer: tuple[CustomGuardrail, bool] | None = None, + endpoint: str = _CONVERSE_STREAM, +) -> bytes: + layers = [(guardrail, holds_until_moderated), *([outer] if outer is not None else [])] + monkeypatch.setattr(litellm, "callbacks", [layer for layer, _ in layers]) + logging_obj = LiteLLMLoggingObj( + model="claude", + messages=[], + stream=True, + call_type="allm_passthrough_route", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + relayed = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_event_stream_iterator_hook( + response=BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(*chunks), endpoint), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + request_data={ + "model": "claude", + "custom_llm_provider": "bedrock", + "endpoint": endpoint, + "litellm_logging_obj": logging_obj, + }, + guardrails=layers, + ) + frames = [] if received is None else received + async for item in relayed: + frame = BedrockPassthroughGuardrailHandler.encode_event_stream_item(item) + if frame: + frames.append(frame) + return b"".join(frames) + + +def _card_stream() -> tuple[bytes, ...]: + return ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "Card: 4111 1111 "}}), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "1111 1111."}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Thanks!"}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + +class _RaisingGuardrail(CustomGuardrail): + def __init__(self, raised: Exception, on_text: str) -> None: + super().__init__(guardrail_name="raiser", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.raised = raised + self.on_text = on_text + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if any(self.on_text in text for text in inputs.get("texts", [])): + raise self.raised + return inputs + + +class _RecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="recorder", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.received_texts: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.received_texts.append(list(inputs.get("texts", []))) + return inputs + + +class _RequestMarksRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="marks", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.received_marks: list[tuple[object, object]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.received_marks.append((request_data.get("stream"), request_data.get("response"))) + return inputs + + +def _two_scan_stream() -> tuple[bytes, ...]: + return ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + *( + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}) + for text in ("Your ", "order ", "ships ", "on ", "Monday.") + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Thanks!"}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + +_NOMA_API_BASE = "https://noma.example.test" + + +def _noma_guardrail(monkeypatch: pytest.MonkeyPatch) -> NomaV2Guardrail: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", None) + return NomaV2Guardrail( + api_key="test-key", + api_base=_NOMA_API_BASE, + guardrail_name="noma", + event_hook="post_call", + default_on=True, + ) + + +class TestConverseStreamFrames: + @pytest.mark.asyncio + async def test_frames_are_cut_at_their_boundaries_whatever_the_transport_chunks(self): + stream = b"".join(_card_stream()) + transport = tuple(stream[i : i + 7] for i in range(0, len(stream), 7)) + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(*transport), _CONVERSE_STREAM + ) + ] + + assert [frame.raw for frame in frames] == [*_card_stream(), b""] + assert [frame.event_type for frame in frames] == [ + "messageStart", + "contentBlockDelta", + "contentBlockDelta", + "contentBlockStop", + "contentBlockDelta", + "contentBlockStop", + "messageStop", + None, + ] + assert [frame.ends_stream for frame in frames] == [False] * 7 + [True] + + @pytest.mark.asyncio + async def test_bytes_that_cannot_be_framed_are_relayed_as_received(self): + frame = _build_event_stream_frame("messageStart", {"role": "assistant"}) + garbage = b"\x00\x00\x00\x05abcdefgh" + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(frame + garbage, b"tail"), _CONVERSE_STREAM + ) + ] + + assert [item.raw for item in frames] == [frame, garbage, b"tail", b""] + assert [item.payload for item in frames[1:]] == [None, None, None] + + @pytest.mark.asyncio + async def test_a_stream_ending_inside_a_frame_relays_the_partial_frame(self): + complete, partial = _card_stream()[:2] + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(complete + partial[:-3]), _CONVERSE_STREAM + ) + ] + + assert [frame.raw for frame in frames] == [complete, partial[:-3], b""] + + @pytest.mark.asyncio + async def test_a_delta_ending_in_half_a_surrogate_pair_is_still_decoded(self): + delta = {"contentBlockIndex": 0, "delta": {"text": "abc\ud83d"}} + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(_build_event_stream_frame("contentBlockDelta", delta)), _CONVERSE_STREAM + ) + ] + + assert frames[0].payload == delta + + @pytest.mark.asyncio + async def test_a_payload_that_is_not_utf8_is_relayed_as_received(self): + frame = _build_frame( + {":event-type": "contentBlockDelta", ":content-type": "application/json", ":message-type": "event"}, + b'{"contentBlockIndex":0,"delta":{"text":"\xff"}}', + ) + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(frame), _CONVERSE_STREAM) + ] + + assert [(item.raw, item.event_type, item.payload) for item in frames] == [ + (frame, "contentBlockDelta", None), + (b"", None, None), + ] + + +def _tool_use_start_frame(index: int, tool_use_id: str, name: str) -> bytes: + tool_use = {"toolUseId": tool_use_id, "name": name, "type": "tool_use"} + return _build_event_stream_frame("contentBlockStart", {"contentBlockIndex": index, "start": {"toolUse": tool_use}}) + + +def _tool_call_stream(tool_input: str) -> tuple[bytes, ...]: + return ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "Checking."}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _tool_use_start_frame(1, "tooluse_1", "lookup_order"), + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"toolUse": {"input": tool_input}}} + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}), + _build_event_stream_frame("messageStop", {"stopReason": "max_tokens"}), + ) + + +def _tool_without_input_stream() -> tuple[bytes, ...]: + return ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + _tool_use_start_frame(0, "tooluse_1", "get_time"), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": ""}}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("messageStop", {"stopReason": "tool_use"}), + ) + + +class TestConverseStreamGuardrailRelay: + @pytest.mark.asyncio + async def test_sampled_stream_relays_frames_as_received_and_scans_one_text_per_block(self, monkeypatch): + guardrail = _CardMaskingGuardrail() + + relayed = await _relay_through_post_call_guardrail(_card_stream(), guardrail, monkeypatch) + + assert relayed == b"".join(_card_stream()) + assert guardrail.received_texts == [[f"Card: {_TEST_CARD}.", "Thanks!"]] + + @pytest.mark.asyncio + async def test_each_sampled_round_scans_the_text_the_caller_received(self, monkeypatch): + guardrail = _CardMaskingGuardrail() + stream = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + *( + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}) + for text in ("Card: 4111 1111 ", "1111 1111.", " More", " text", " here") + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + relayed = await _relay_through_post_call_guardrail(stream, guardrail, monkeypatch) + + assert relayed == b"".join(stream) + assert guardrail.received_texts == [ + [f"Card: {_TEST_CARD}. More text"], + [f"Card: {_TEST_CARD}. More text here"], + ] + + @pytest.mark.asyncio + async def test_stream_held_until_moderated_applies_masks_per_block(self, monkeypatch): + guardrail = _CardMaskingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + _card_stream(), guardrail, monkeypatch, holds_until_moderated=True + ) + + frames = _decoded_frames(relayed) + assert [kind for _, kind, _ in frames] == [ + "messageStart", + "contentBlockDelta", + "contentBlockDelta", + "contentBlockStop", + "contentBlockDelta", + "contentBlockStop", + "messageStop", + ] + assert TestDeAnonymizeConverseStream._deltas_by_block(relayed, "text") == {0: "Card: ****.", 1: "Thanks!"} + assert relayed.endswith(b"".join(_card_stream()[3:])) + + @pytest.mark.parametrize( + "ending", + [ + pytest.param( + BedrockPassthroughGuardrailHandler.event_stream_error_item("Model stream error", 500), + id="exception-frame", + ), + pytest.param( + _build_frame( + {":message-type": "error", ":error-code": "InternalFailure", ":error-message": "boom"}, b"" + ), + id="error-frame", + ), + pytest.param(b"", id="upstream-closed-without-messageStop"), + ], + ) + @pytest.mark.asyncio + async def test_held_stream_that_ends_without_messageStop_still_applies_masks(self, monkeypatch, ending: bytes): + stream = (*_card_stream()[:3], ending) + + relayed = await _relay_through_post_call_guardrail( + stream, _CardMaskingGuardrail(), monkeypatch, holds_until_moderated=True + ) + + assert TestDeAnonymizeConverseStream._deltas_by_block(relayed, "text") == {0: "Card: ****."} + assert relayed.endswith(ending) + + @pytest.mark.asyncio + async def test_a_sampled_guardrail_rewriting_at_the_end_never_reaches_frames_a_later_guardrail_holds( + self, monkeypatch + ): + stream = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + *( + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}) + for text in ("Card: 4111", " 1111", " 1111", " 1111", " now.") + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + held = _RecordingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + stream, _CardMaskingGuardrail(), monkeypatch, outer=(held, True) + ) + + assert relayed == b"".join(stream) + assert held.received_texts[-1] == [f"Card: {_TEST_CARD} now."] + + @pytest.mark.parametrize("filler_frames", range(8)) + @pytest.mark.asyncio + async def test_a_held_guardrail_still_masks_after_a_sampled_guardrail_ends_the_stream_with_its_message( + self, monkeypatch, filler_frames: int + ): + from litellm.exceptions import ModifyResponseException + + sampled = _RaisingGuardrail( + ModifyResponseException(message="Sorry, I can't share that.", model="claude", request_data={}), + "FORBIDDEN", + ) + stream = ( + *_card_stream()[:3], + *( + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": " ok"}}) + for _ in range(filler_frames) + ), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": " FORBIDDEN"}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + relayed = await _relay_through_post_call_guardrail( + stream, sampled, monkeypatch, outer=(_CardMaskingGuardrail(), True) + ) + + texts = TestDeAnonymizeConverseStream._deltas_by_block(relayed, "text") + assert texts[0].startswith("Card: ****."), texts + assert texts[max(texts)] == "Sorry, I can't share that.", texts + assert _decoded_frames(relayed)[-1] == ("event", "messageStop", {"stopReason": "guardrail_intervened"}) + + @pytest.mark.asyncio + async def test_noma_scans_carry_stream_and_only_the_last_one_the_finished_answer( + self, monkeypatch, respx_mock: respx.MockRouter + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + + await _relay_through_post_call_guardrail(_two_scan_stream(), _noma_guardrail(monkeypatch), monkeypatch) + + sent = [json.loads(call.request.content)["request_data"] for call in route.calls] + answer = [{"text": "Your order ships on Monday."}, {"text": "Thanks!"}] + finished = {"output": {"message": {"role": "assistant", "content": answer}}, "stopReason": "end_turn"} + assert [(scan.get("stream"), scan.get("response")) for scan in sent] == [(True, None), (True, finished)] + + @pytest.mark.asyncio + async def test_noma_gets_the_finished_answer_when_the_last_sampled_scan_already_covered_it( + self, monkeypatch, respx_mock: respx.MockRouter + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + later = _RequestMarksRecordingGuardrail() + stream = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + *( + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}) + for text in ("Your ", "order ", "ships ", "Monday.") + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + await _relay_through_post_call_guardrail( + stream, _noma_guardrail(monkeypatch), monkeypatch, outer=(later, False) + ) + + sent = [json.loads(call.request.content)["request_data"] for call in route.calls] + finished = { + "output": {"message": {"role": "assistant", "content": [{"text": "Your order ships Monday."}]}}, + "stopReason": "end_turn", + } + assert [(scan.get("stream"), scan.get("response")) for scan in sent] == [(True, None), (True, finished)] + assert later.received_marks == [(None, None)] + + @pytest.mark.asyncio + async def test_noma_gets_the_finished_answer_with_its_reasoning_and_tool_calls_as_converse_blocks( + self, monkeypatch, respx_mock: respx.MockRouter + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + stream = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": "Look it up."}}} + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Checking."}}), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}), + _tool_use_start_frame(2, "tooluse_1", "lookup_order"), + *( + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 2, "delta": {"toolUse": {"input": fragment}}} + ) + for fragment in ('{"order_', 'id": "A1"}') + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 2}), + _tool_use_start_frame(3, "tooluse_2", "list_orders"), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 3}), + _build_event_stream_frame("messageStop", {"stopReason": "tool_use"}), + ) + + await _relay_through_post_call_guardrail(stream, _noma_guardrail(monkeypatch), monkeypatch) + + finished = json.loads(route.calls[-1].request.content)["request_data"]["response"] + assert finished == { + "output": { + "message": { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "Look it up."}}}, + {"text": "Checking."}, + { + "toolUse": { + "toolUseId": "tooluse_1", + "name": "lookup_order", + "type": "tool_use", + "input": {"order_id": "A1"}, + } + }, + {"toolUse": {"toolUseId": "tooluse_2", "name": "list_orders", "type": "tool_use", "input": {}}}, + ], + } + }, + "stopReason": "tool_use", + } + + @pytest.mark.parametrize( + "tool_input", + [ + pytest.param('{"order_id": "A', id="cut-short"), + pytest.param('{"order_id": NaN}', id="nan"), + pytest.param("[" * 10_000, id="nested-too-deep"), + ], + ) + @pytest.mark.asyncio + async def test_noma_gets_a_tool_input_that_is_not_json_as_streamed( + self, monkeypatch, respx_mock: respx.MockRouter, tool_input: str + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + + relayed = await _relay_through_post_call_guardrail( + _tool_call_stream(tool_input), _noma_guardrail(monkeypatch), monkeypatch + ) + + assert relayed == b"".join(_tool_call_stream(tool_input)) + finished = json.loads(route.calls[-1].request.content)["request_data"]["response"] + assert finished["output"]["message"]["content"][1]["toolUse"]["input"] == tool_input + + @pytest.mark.parametrize("held", [False, True], ids=["sampled", "held"]) + @pytest.mark.asyncio + async def test_noma_alone_gets_one_ended_scan_of_a_reply_that_only_calls_a_tool_without_input( + self, monkeypatch, respx_mock: respx.MockRouter, held: bool + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + later = _RecordingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + _tool_without_input_stream(), + _noma_guardrail(monkeypatch), + monkeypatch, + holds_until_moderated=held, + outer=(later, False), + ) + + assert relayed == b"".join(_tool_without_input_stream()) + scans = [json.loads(call.request.content) for call in route.calls] + tool_call = {"toolUse": {"toolUseId": "tooluse_1", "name": "get_time", "type": "tool_use", "input": {}}} + finished = {"output": {"message": {"role": "assistant", "content": [tool_call]}}, "stopReason": "tool_use"} + assert [ + ( + scan["input_type"], + scan["inputs"]["texts"], + scan["request_data"]["stream"], + scan["request_data"]["response"], + ) + for scan in scans + ] == [("response", [], True, finished)] + assert later.received_texts == [] + + @pytest.mark.asyncio + async def test_noma_gets_no_scan_of_a_reply_with_no_text_and_no_tool_call( + self, monkeypatch, respx_mock: respx.MockRouter + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + stream = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + relayed = await _relay_through_post_call_guardrail(stream, _noma_guardrail(monkeypatch), monkeypatch) + + assert relayed == b"".join(stream) + assert route.calls.call_count == 0 + + @pytest.mark.asyncio + async def test_noma_gets_no_scan_of_a_tool_call_without_input_before_the_reply_ends( + self, monkeypatch, respx_mock: respx.MockRouter + ): + route = respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(*_tool_without_input_stream()), _CONVERSE_STREAM + ) + ] + + await BedrockPassthroughGuardrailHandler().process_output_streaming_response( + frames[:4], + _noma_guardrail(monkeypatch), + request_data={"model": "claude", "endpoint": "model/claude/converse-stream"}, + ) + + assert route.calls.call_count == 0 + + @pytest.mark.asyncio + async def test_guardrails_scanning_after_noma_see_the_request_without_its_stream_marks( + self, monkeypatch, respx_mock: respx.MockRouter + ): + respx_mock.post(f"{_NOMA_API_BASE}/litellm/guardrail").respond(json={"action": "NONE"}) + later = _RequestMarksRecordingGuardrail() + + await _relay_through_post_call_guardrail( + _two_scan_stream(), _noma_guardrail(monkeypatch), monkeypatch, outer=(later, False) + ) + + assert later.received_marks == [(None, None), (None, None)] + + @pytest.mark.asyncio + async def test_block_after_frames_were_relayed_raises_for_the_relay_to_report(self, monkeypatch): + from fastapi import HTTPException + + block = HTTPException(status_code=400, detail="Violated guardrail policy") + relayed: list[bytes] = [] + + with pytest.raises(HTTPException) as raised: + await _relay_through_post_call_guardrail( + _card_stream(), _RaisingGuardrail(block, "Thanks"), monkeypatch, received=relayed + ) + + assert raised.value is block + assert relayed == list(_card_stream()[:4]) + + @pytest.mark.asyncio + async def test_guardrail_message_ends_the_message_after_closing_the_open_block(self, monkeypatch): + from litellm.exceptions import ModifyResponseException + + guardrail = _RaisingGuardrail( + ModifyResponseException(message="Sorry, I can't share that.", model="claude", request_data={}), + "1111 1111.", + ) + stream = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}), + *( + _build_event_stream_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}) + for text in ("Card: ", "4111 1111 ", "1111 1111.", " Done") + ), + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}), + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}), + ) + + relayed = await _relay_through_post_call_guardrail(stream, guardrail, monkeypatch) + + assert relayed.startswith(b"".join(stream[:4])) + assert _decoded_frames(relayed)[4:] == [ + ("event", "contentBlockStop", {"contentBlockIndex": 0}), + ("event", "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Sorry, I can't share that."}}), + ("event", "contentBlockStop", {"contentBlockIndex": 1}), + ("event", "messageStop", {"stopReason": "guardrail_intervened"}), + ] + + @pytest.mark.asyncio + async def test_guardrail_message_on_a_held_stream_is_a_whole_message(self, monkeypatch): + from litellm.exceptions import ModifyResponseException + + guardrail = _RaisingGuardrail( + ModifyResponseException( + message="Sorry, I can't share that.", + model="claude", + request_data={"custom_llm_provider": "bedrock"}, + ), + "Card", + ) + + relayed = await _relay_through_post_call_guardrail( + _card_stream(), guardrail, monkeypatch, holds_until_moderated=True + ) + + assert _decoded_frames(relayed) == [ + ("event", "messageStart", {"role": "assistant"}), + ("event", "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "Sorry, I can't share that."}}), + ("event", "contentBlockStop", {"contentBlockIndex": 0}), + ("event", "messageStop", {"stopReason": "guardrail_intervened"}), + ] + + @pytest.mark.parametrize( + ("status_code", "exception_type"), + [(400, "validationException"), (503, "internalServerException")], + ) + def test_stream_error_frame_type_follows_the_status(self, status_code: int, exception_type: str): + item = BedrockPassthroughGuardrailHandler.event_stream_error_item("scan failed", status_code) + + assert _decoded_frames(item) == [("exception", exception_type, {"message": "scan failed"})] + + +class _InvokeTextBlock(TypedDict): + type: ReadOnly[Literal["text"]] + text: ReadOnly[str] + + +class _InvokeToolUseBlock(TypedDict): + type: ReadOnly[Literal["tool_use"]] + id: ReadOnly[str] + name: ReadOnly[str] + input: ReadOnly[Mapping[str, JsonValue]] + + +class _InvokeUsage(TypedDict, total=False): + input_tokens: ReadOnly[int] + output_tokens: ReadOnly[int] + + +class _InvokeMessage(TypedDict, total=False): + id: ReadOnly[str] + type: ReadOnly[Literal["message"]] + role: ReadOnly[Literal["assistant"]] + content: ReadOnly[Sequence[_InvokeTextBlock | _InvokeToolUseBlock]] + model: ReadOnly[str] + stop_reason: ReadOnly[str | None] + usage: ReadOnly[_InvokeUsage] + + +class _InvokeTextDelta(TypedDict): + type: ReadOnly[Literal["text_delta"]] + text: ReadOnly[str] + + +class _InvokeJsonDelta(TypedDict): + type: ReadOnly[Literal["input_json_delta"]] + partial_json: ReadOnly[str] + + +class _InvokeThinkingDelta(TypedDict): + type: ReadOnly[Literal["thinking_delta"]] + thinking: ReadOnly[str] + + +class _InvokeSignatureDelta(TypedDict): + type: ReadOnly[Literal["signature_delta"]] + signature: ReadOnly[str] + + +_InvokeDelta: TypeAlias = _InvokeTextDelta | _InvokeJsonDelta | _InvokeThinkingDelta | _InvokeSignatureDelta + + +class _InvokeMessageStart(TypedDict): + type: ReadOnly[Literal["message_start"]] + message: ReadOnly[_InvokeMessage] + + +class _InvokeBlockStart(TypedDict): + type: ReadOnly[Literal["content_block_start"]] + index: ReadOnly[int] + content_block: ReadOnly[_InvokeTextBlock | _InvokeToolUseBlock] + + +class _InvokeBlockDelta(TypedDict): + type: ReadOnly[Literal["content_block_delta"]] + index: ReadOnly[int] + delta: ReadOnly[_InvokeDelta] + + +class _InvokeBlockStop(TypedDict): + type: ReadOnly[Literal["content_block_stop"]] + index: ReadOnly[int] + + +class _InvokeStopReason(TypedDict): + stop_reason: ReadOnly[str] + + +class _InvokeMessageDelta(TypedDict): + type: ReadOnly[Literal["message_delta"]] + delta: ReadOnly[_InvokeStopReason] + usage: ReadOnly[_InvokeUsage] + + +class _InvokeMessageStop(TypedDict): + type: ReadOnly[Literal["message_stop"]] + + +_InvokeEvent: TypeAlias = ( + _InvokeMessageStart + | _InvokeBlockStart + | _InvokeBlockDelta + | _InvokeBlockStop + | _InvokeMessageDelta + | _InvokeMessageStop +) + + +def _invoke_raw_chunk_frame(event_json: bytes) -> bytes: + import base64 + + return _build_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, + json.dumps({"bytes": base64.b64encode(event_json).decode(), "p": "abcdef"}).encode(), + ) + + +def _invoke_chunk_frame(event: _InvokeEvent) -> bytes: + return _invoke_raw_chunk_frame(json.dumps(event).encode()) + + +def _invoke_stream(*deltas: str) -> tuple[bytes, ...]: + return ( + _invoke_chunk_frame( + { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "content": [], "model": "claude"}, + } + ), + _invoke_chunk_frame({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _invoke_chunk_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": t}}) + for t in deltas + ), + _invoke_chunk_frame({"type": "content_block_stop", "index": 0}), + _invoke_chunk_frame( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}} + ), + _invoke_chunk_frame({"type": "message_stop"}), + ) + + +def _invoke_card_stream() -> tuple[bytes, ...]: + return _invoke_stream("Card: 4111 1111 ", "1111 1111.", " Thanks!") + + +def _decoded_invoke_frames(raw: bytes) -> list[tuple[str, dict[str, JsonValue]]]: + import base64 + import json + + return [ + (message_type, json.loads(base64.b64decode(payload["bytes"])) if kind == "chunk" else payload) + for message_type, kind, payload in _decoded_frames(raw) + ] + + +def _delta_text(event: Mapping[str, JsonValue]) -> str: + delta = event["delta"] + assert isinstance(delta, dict) and delta.get("type") == "text_delta", event + text = delta["text"] + assert isinstance(text, str), event + return text + + +def _streamed_text(raw: bytes) -> str: + return "".join( + _delta_text(event) + for _, event in _decoded_invoke_frames(raw) + if event.get("type") == "content_block_delta" + and isinstance(delta := event.get("delta"), dict) + and delta.get("type") == "text_delta" + ) + + +class _RequestDataRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="request-recorder", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.scans: list[tuple[object, int, object]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + response = request_data.get("response") + self.scans.append( + ( + request_data.get("stream"), + len(request_data.get("responses") or ()), + response.choices[0].finish_reason if response is not None else None, + ) + ) + return inputs + + +class _ToolInputRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="tool-recorder", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.tool_arguments: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.tool_arguments.append([call.function.arguments for call in inputs.get("tool_calls") or []]) + return inputs + + +def _invoke_delta_frame(index: int, delta: _InvokeDelta) -> bytes: + return _invoke_chunk_frame({"type": "content_block_delta", "index": index, "delta": delta}) + + +class TestBufferedInvokeStream: + @pytest.mark.asyncio + async def test_masks_are_written_back_into_the_chunk_frames_that_carried_the_text( + self, monkeypatch: pytest.MonkeyPatch + ): + guardrail = _CardMaskingGuardrail() + stream = b"".join(_invoke_card_stream()) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["Card: 4111 1111 1111 1111. Thanks!"]] + assert _streamed_text(result) == "Card: ****. Thanks!" + assert [event for _, event in _decoded_invoke_frames(result) if event["type"] != "content_block_delta"] == [ + event for _, event in _decoded_invoke_frames(stream) if event["type"] != "content_block_delta" + ] + + @pytest.mark.asyncio + async def test_each_content_block_is_scanned_and_masked_on_its_own(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + stream = b"".join( + ( + _invoke_delta_frame(0, {"type": "thinking_delta", "thinking": "The card is 4111 1111 1111 1111"}), + _invoke_delta_frame(0, {"type": "signature_delta", "signature": "sig-4111 1111 1111 1111"}), + _invoke_delta_frame(1, {"type": "text_delta", "text": "Paying with 4111 1111 "}), + _invoke_delta_frame(1, {"type": "text_delta", "text": "1111 1111 now"}), + _invoke_delta_frame(2, {"type": "input_json_delta", "partial_json": '{"card": "4111 1111 1111 1111"}'}), + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + events = [event for _, event in _decoded_invoke_frames(result) if event["type"] == "content_block_delta"] + assert events[0]["delta"] == {"type": "thinking_delta", "thinking": "The card is ****"} + assert events[1]["delta"] == {"type": "signature_delta", "signature": "sig-4111 1111 1111 1111"} + assert _delta_text(events[2]) + _delta_text(events[3]) == "Paying with **** now" + assert events[4]["delta"] == {"type": "input_json_delta", "partial_json": '{"card": "****"}'} + + @pytest.mark.asyncio + async def test_two_text_blocks_reach_the_guardrail_as_two_texts(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + stream = b"".join( + ( + _invoke_delta_frame(0, {"type": "text_delta", "text": "First 4111 1111 1111 1111"}), + _invoke_delta_frame(1, {"type": "text_delta", "text": "Second 4111 1111 1111 1111"}), + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["First 4111 1111 1111 1111", "Second 4111 1111 1111 1111"]] + events = [event for _, event in _decoded_invoke_frames(result)] + assert [_delta_text(event) for event in events[:2]] == ["First ****", "Second ****"] + + @pytest.mark.asyncio + async def test_a_delta_whose_text_is_not_a_string_is_left_as_received(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + odd = _invoke_raw_chunk_frame( + b'{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":4111}}' + ) + stream = b"".join( + ( + _invoke_delta_frame(0, {"type": "text_delta", "text": "Card 4111 1111 1111 1111"}), + odd, + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["Card 4111 1111 1111 1111"]] + frames = _decoded_invoke_frames(result) + assert _delta_text(frames[0][1]) == "Card ****" + assert odd in result + + @pytest.mark.asyncio + async def test_a_stream_without_text_skips_the_post_call_hooks(self): + hook = AsyncMock() + stream = b"".join( + ( + _invoke_chunk_frame({"type": "message_start", "message": {"id": "msg_1", "content": []}}), + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( + body_bytes=stream, + proxy_logging_obj=MagicMock(post_call_success_hook=hook), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + data={"custom_llm_provider": "bedrock", "endpoint": _INVOKE_STREAM}, + ) + + hook.assert_not_awaited() + assert result == stream + + @pytest.mark.asyncio + async def test_a_stream_the_guardrail_leaves_alone_is_returned_byte_for_byte(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _RecordingGuardrail() + stream = b"".join(_invoke_stream("Hello ", "there")) + b"trailing" + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["Hello there"]] + assert result == stream + + @pytest.mark.asyncio + async def test_a_block_raises_before_any_frame_is_returned(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _RaisingGuardrail(GuardrailBlocked("Blocked"), "Card") + + with pytest.raises(GuardrailBlocked): + await _run_post_call_guardrail_on_stream( + b"".join(_invoke_card_stream()), guardrail, monkeypatch, endpoint=_INVOKE_STREAM + ) + + +class TestInvokeStreamGuardrailRelay: + @pytest.mark.asyncio + async def test_frames_are_cut_at_their_boundaries_whatever_the_transport_chunks(self): + stream = b"".join(_invoke_card_stream()) + transport = tuple(stream[i : i + 7] for i in range(0, len(stream), 7)) + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(*transport), _INVOKE_STREAM + ) + ] + + assert [frame.raw for frame in frames] == [*_invoke_card_stream(), b""] + assert [frame.event["type"] for frame in frames[:-1]] == [ + event["type"] for _, event in _decoded_invoke_frames(stream) + ] + assert [BedrockPassthroughGuardrailHandler.encode_event_stream_item(frame) for frame in frames] == [ + *_invoke_card_stream(), + b"", + ] + + @pytest.mark.parametrize( + "payload", + [ + pytest.param(b"not json", id="payload-not-json"), + pytest.param(b'{"bytes": "%%%"}', id="bytes-not-base64"), + pytest.param(b'{"bytes": "bm90IGpzb24="}', id="decoded-not-json"), + pytest.param(b'{"bytes": "eyJpbmRleCI6IDB9"}', id="event-without-type"), + pytest.param(b'{"p": "abc"}', id="no-bytes"), + pytest.param(b'{"bytes": 5}', id="bytes-not-a-string"), + ], + ) + @pytest.mark.asyncio + async def test_a_chunk_that_does_not_decode_is_relayed_as_received(self, payload: bytes): + frame = _build_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, payload + ) + + frames = [ + item async for item in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(frame), _INVOKE_STREAM) + ] + + assert [(item.raw, item.event) for item in frames] == [(frame, None), (b"", None)] + assert BedrockPassthroughGuardrailHandler.encode_event_stream_item(frames[0]) == frame + + @pytest.mark.asyncio + async def test_sampled_stream_relays_frames_as_received_and_scans_the_text(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + _invoke_card_stream(), guardrail, monkeypatch, endpoint=_INVOKE_STREAM + ) + + assert relayed == b"".join(_invoke_card_stream()) + assert guardrail.received_texts == [[f"Card: {_TEST_CARD}. Thanks!"]] + + @pytest.mark.asyncio + async def test_stream_held_until_moderated_writes_the_masks_back_into_the_chunk_frames( + self, monkeypatch: pytest.MonkeyPatch + ): + stream = _invoke_card_stream() + + relayed = await _relay_through_post_call_guardrail( + stream, _CardMaskingGuardrail(), monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert [event["type"] for _, event in _decoded_invoke_frames(relayed)] == [ + event["type"] for _, event in _decoded_invoke_frames(b"".join(stream)) + ] + assert _streamed_text(relayed) == "Card: ****. Thanks!" + assert {payload.get("p") for _, _, payload in _decoded_frames(relayed)} == {"abcdef"} + assert relayed.startswith(b"".join(stream[:2])) and relayed.endswith(b"".join(stream[5:])) + + @pytest.mark.parametrize( + "ending", + [ + pytest.param( + BedrockPassthroughGuardrailHandler.event_stream_error_item("Model stream error", 500), + id="exception-frame", + ), + pytest.param(b"", id="upstream-closed-without-a-stop-reason"), + ], + ) + @pytest.mark.asyncio + async def test_held_stream_that_ends_without_a_stop_reason_still_applies_masks( + self, monkeypatch: pytest.MonkeyPatch, ending: bytes + ): + stream = (*_invoke_card_stream()[:4], ending) + + relayed = await _relay_through_post_call_guardrail( + stream, _CardMaskingGuardrail(), monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert _streamed_text(relayed) == "Card: ****." + assert relayed.endswith(ending) + + @pytest.mark.asyncio + async def test_a_later_guardrail_scans_the_masks_a_held_guardrail_wrote_back(self, monkeypatch: pytest.MonkeyPatch): + held = _RecordingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + _invoke_card_stream(), + _CardMaskingGuardrail(), + monkeypatch, + holds_until_moderated=True, + outer=(held, True), + endpoint=_INVOKE_STREAM, + ) + + assert _streamed_text(relayed) == "Card: ****. Thanks!" + assert held.received_texts == [["Card: ****. Thanks!"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_block_withholds_the_frames_after_the_stop_reason( + self, monkeypatch: pytest.MonkeyPatch + ): + from fastapi import HTTPException + + guardrail = _RaisingGuardrail(HTTPException(status_code=400, detail="Violated guardrail policy"), "Thanks") + guardrail.streaming_end_of_stream_only = True + relayed: list[bytes] = [] + + with pytest.raises(HTTPException): + await _relay_through_post_call_guardrail( + _invoke_card_stream(), guardrail, monkeypatch, received=relayed, endpoint=_INVOKE_STREAM + ) + + assert relayed == list(_invoke_card_stream()[:-2]) + + @pytest.mark.asyncio + async def test_block_after_frames_were_relayed_raises_for_the_relay_to_report( + self, monkeypatch: pytest.MonkeyPatch + ): + from fastapi import HTTPException + + block = HTTPException(status_code=400, detail="Violated guardrail policy") + relayed: list[bytes] = [] + + with pytest.raises(HTTPException) as raised: + await _relay_through_post_call_guardrail( + _invoke_card_stream(), + _RaisingGuardrail(block, "Thanks"), + monkeypatch, + received=relayed, + endpoint=_INVOKE_STREAM, + ) + + assert raised.value is block + assert relayed == list(_invoke_card_stream()[:4]) + + @pytest.mark.asyncio + async def test_guardrail_message_ends_the_message_after_closing_the_open_block( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.exceptions import ModifyResponseException + + guardrail = _RaisingGuardrail( + ModifyResponseException(message="Sorry, I can't share that.", model="claude", request_data={}), + "1111 1111.", + ) + stream = _invoke_stream("Card: ", "4111 1111 ", "1111 1111.", " Done") + + relayed = await _relay_through_post_call_guardrail(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert relayed.startswith(b"".join(stream[:4])) + assert [event for _, event in _decoded_invoke_frames(relayed[len(b"".join(stream[:4])) :])] == [ + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "Sorry, I can't share that."}, + }, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 0}, + }, + {"type": "message_stop"}, + ] + + @pytest.mark.asyncio + async def test_guardrail_message_on_a_held_stream_is_a_whole_message(self, monkeypatch: pytest.MonkeyPatch): + from litellm.exceptions import ModifyResponseException + + guardrail = _RaisingGuardrail( + ModifyResponseException( + message="Sorry, I can't share that.", + model="claude", + request_data={"custom_llm_provider": "bedrock", "endpoint": _INVOKE_STREAM}, + ), + "Card", + ) + + relayed = await _relay_through_post_call_guardrail( + _invoke_card_stream(), guardrail, monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert [event["type"] for _, event in _decoded_invoke_frames(relayed)] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert _streamed_text(relayed) == "Sorry, I can't share that." + + @pytest.mark.parametrize( + ("endpoint", "model", "body", "relayable"), + [ + ("model/x/converse-stream", "meta.llama4-maverick-17b-instruct-v1:0", {}, True), + (_INVOKE_STREAM, "anthropic.claude-sonnet-5-v1:0", {}, True), + (_INVOKE_STREAM, "us.anthropic.claude-sonnet-5-v1:0", {}, True), + (_INVOKE_STREAM, "invoke/anthropic.claude-sonnet-5-v1:0", {}, True), + (_INVOKE_STREAM, "meta.llama4-maverick-17b-instruct-v1:0", {"prompt": "hi"}, False), + (_INVOKE_STREAM, _APPLICATION_PROFILE, {"anthropic_version": "bedrock-2023-05-31"}, True), + (_INVOKE_STREAM, _APPLICATION_PROFILE, {"prompt": "hi"}, False), + (_INVOKE_STREAM, _APPLICATION_PROFILE, None, False), + (_INVOKE_STREAM, _APPLICATION_PROFILE, ["anthropic_version"], False), + ("model/anthropic.claude-sonnet-5-v1:0/invoke", "anthropic.claude-sonnet-5-v1:0", {}, False), + ("model/x/invoke", _APPLICATION_PROFILE, {"anthropic_version": "bedrock-2023-05-31"}, False), + ], + ) + def test_relayed_and_buffered_streams_are_converse_and_claude_invoke_streams( + self, endpoint: str, model: str, body: object, relayable: bool + ): + assert BedrockPassthroughGuardrailHandler.event_stream_endpoint_is_relayable(endpoint, model, body) is relayable + assert ( + BedrockPassthroughGuardrailHandler.event_stream_endpoint_is_de_anonymizable(endpoint, model, body) + is relayable + ) + + @pytest.mark.asyncio + async def test_each_sampled_scan_sees_a_streaming_request_and_every_event_relayed_so_far( + self, monkeypatch: pytest.MonkeyPatch + ): + guardrail = _RequestDataRecordingGuardrail() + guardrail.streaming_sampling_rate = 2 + + await _relay_through_post_call_guardrail( + _invoke_stream("Card: ", "4111 1111 ", "1111 1111", "."), guardrail, monkeypatch, endpoint=_INVOKE_STREAM + ) + + assert len(guardrail.scans) > 1 + assert all(stream is True for stream, _, _ in guardrail.scans) + responses_seen = [responses for _, responses, _ in guardrail.scans] + assert responses_seen == sorted(set(responses_seen)), responses_seen + assert [finish_reason for _, _, finish_reason in guardrail.scans[:-1]] == [None] * (len(guardrail.scans) - 1) + + @pytest.mark.asyncio + async def test_the_held_scan_sees_the_whole_stream_and_its_finish_reason(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _RequestDataRecordingGuardrail() + stream = ( + _invoke_chunk_frame( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "model": "claude", + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + }, + } + ), + *_invoke_card_stream()[1:], + ) + + await _relay_through_post_call_guardrail( + stream, guardrail, monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert guardrail.scans == [(True, len(stream) + 1, "stop")] + + @pytest.mark.parametrize("holds_until_moderated", [True, False], ids=["held", "sampled"]) + @pytest.mark.asyncio + async def test_a_stream_the_upstream_breaks_off_still_gets_its_tool_input_scanned( + self, monkeypatch: pytest.MonkeyPatch, holds_until_moderated: bool + ): + guardrail = _ToolInputRecordingGuardrail() + stream = ( + _invoke_chunk_frame( + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + ), + _invoke_delta_frame(0, {"type": "text_delta", "text": "Calling the tool"}), + _invoke_chunk_frame({"type": "content_block_stop", "index": 0}), + _invoke_chunk_frame( + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}}, + } + ), + _invoke_delta_frame(1, {"type": "input_json_delta", "partial_json": '{"card": "4111 1111 1111 1111"}'}), + _invoke_chunk_frame({"type": "content_block_stop", "index": 1}), + _build_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + b'{"message":"Model stream error"}', + ), + ) + + await _relay_through_post_call_guardrail( + stream, guardrail, monkeypatch, holds_until_moderated=holds_until_moderated, endpoint=_INVOKE_STREAM + ) + + assert guardrail.tool_arguments[-1] == ['{"card": "4111 1111 1111 1111"}'] + + @pytest.mark.asyncio + async def test_sampled_scan_key_reads_the_received_text_once_a_guardrail_rewrote_it(self): + delta: _InvokeBlockDelta = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "In Paris"}, + } + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(_invoke_chunk_frame(delta)), _INVOKE_STREAM + ) + ] + received_key = BedrockPassthroughGuardrailHandler().get_streaming_scan_key(frames) + frames[0].guardrailed_event = {**delta, "delta": {"type": "text_delta", "text": "In [CITY]"}} + + rewritten_key = BedrockPassthroughGuardrailHandler().get_streaming_scan_key(frames) + + assert received_key is not None and received_key.texts == ("In Paris",) + assert rewritten_key == received_key diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 936528f93e7..da5c4699925 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -17,6 +17,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.passthrough.transformation import ( BedrockPassthroughConfig, is_bedrock_streaming_endpoint, @@ -691,6 +692,58 @@ def test_converse_stream_collector_keeps_usage_without_retaining_the_stream(): ) +def test_converse_stream_cut_before_usage_estimates_prompt_tokens_from_the_request(): + logging_obj: Final = _converse_stream_logging_obj() + logging_obj.model_call_details["input"] = { + "system": [{"text": "You are a terse assistant."}], + "messages": [ + {"role": "user", "content": [{"text": "Write a story about a lighthouse keeper."}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "lookup", "input": {"query": "harbor log"}}}], + }, + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [{"text": "calm water"}, {"json": {"wind": "steady"}}], + } + } + ], + }, + ], + } + cut: Final = _converse_stream_collector() + _feed(cut, _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, ["Once ", "upon"])) + response: Final = cut.build_logged_response(logging_obj) + assert isinstance(response, ModelResponse) + expected_messages: Final = [ + {"role": "system", "content": "You are a terse assistant."}, + {"role": "user", "content": "Write a story about a lighthouse keeper."}, + {"role": "user", "content": '{"query": "harbor log"}'}, + {"role": "user", "content": 'calm water\n{"wind": "steady"}'}, + ] + assert response.usage.prompt_tokens == litellm.token_counter(model=logging_obj.model, messages=expected_messages) + assert response.usage.prompt_tokens > 0 + + for request_without_turns in ("not a Converse request", {"prompt": "a model-native invoke body"}): + logging_obj.model_call_details["input"] = request_without_turns + unestimated = cut.build_logged_response(logging_obj) + assert isinstance(unestimated, ModelResponse) + assert (unestimated.usage.prompt_tokens, unestimated.usage.completion_tokens) == ( + 0, + response.usage.completion_tokens, + ), request_without_turns + + full: Final = _converse_stream_collector() + _feed(full, _converse_text_stream(["Once ", "upon"])) + full_response: Final = full.build_logged_response(logging_obj) + assert isinstance(full_response, ModelResponse) + assert full_response.usage.prompt_tokens == 25 + + class _CaptureAsyncSuccessCallback(CustomLogger): def __init__(self) -> None: super().__init__() @@ -734,6 +787,32 @@ async def test_bedrock_stream_passthrough_logs_cost_and_standard_metadata(endpoi assert standard_logging_object["model"] == f"bedrock/{CONVERSE_MODEL}" +@pytest.mark.asyncio +async def test_bedrock_stream_flush_says_whether_it_logged_the_stream() -> None: + stream: Final = _converse_text_stream(["Converse response"]) + part_of_the_first_frame: Final = _stream_collector(CONVERSE_STREAM_ENDPOINT) + _feed(part_of_the_first_frame, stream[:12]) + whole_stream: Final = _stream_collector(CONVERSE_STREAM_ENDPOINT) + _feed(whole_stream, stream) + nothing_logged: Final = _CaptureAsyncSuccessCallback() + logged: Final = _CaptureAsyncSuccessCallback() + + assert ( + await _stream_logging_obj(CONVERSE_STREAM_ENDPOINT, [nothing_logged]).async_flush_passthrough_collected_chunks( + collector=part_of_the_first_frame + ) + is False + ) + assert nothing_logged.kwargs is None + assert ( + await _stream_logging_obj(CONVERSE_STREAM_ENDPOINT, [logged]).async_flush_passthrough_collected_chunks( + collector=whole_stream + ) + is True + ) + assert logged.kwargs is not None + + def test_converse_stream_collector_keeps_tool_calls_between_text_runs(): stream = ( _event_frame("messageStart", {"role": "assistant"}) @@ -820,3 +899,293 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): collector.add(_event_frame("messageStart", {"role": "assistant"})) assert collector.build_logged_response(_converse_stream_logging_obj()) is None + + +REQUEST_METADATA_HEADER = "x-amzn-bedrock-request-metadata" +REQUEST_METADATA_FIELDS = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"] +PROXY_IDENTITY = { + "user_api_key_alias": "prod-key", + "user_api_key_team_alias": "platform", + "spend_logs_metadata": {"agent": "triage-bot"}, +} +EXPECTED_REQUEST_METADATA = { + "user_api_key_alias": "prod-key", + "user_api_key_team_alias": "platform", + "agent": "triage-bot", +} +CONVERSE_BODY = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} +INVOKE_BODY = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} + + +async def _sent_passthrough_request( + action: str, body: dict[str, object], identity: dict[str, object], **kwargs: object +) -> httpx.Request: + requests: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"ok": True}, request=request) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-pt", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAROUTERKEY", + "aws_secret_access_key": "router-secret", + "aws_region_name": "us-west-2", + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-pt", + method="POST", + endpoint=f"/model/bedrock-pt/{action}", + data=body, + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + litellm_metadata=identity, + **kwargs, + ) + return requests[0] + + +def _signed_header_names(request: httpx.Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "query_params"), + [("converse", None), ("converse-stream", None), ("converse/", None), ("converse", {"trace": "1"})], +) +async def test_converse_passthrough_signs_proxy_identity_into_request_metadata( + action: str, query_params: dict[str, str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}} + + request: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, request_query_params=query_params + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"}, + } + assert caller_body["requestMetadata"] == {"user_api_key_team_alias": "forged", "ticket": "T-42"} + assert REQUEST_METADATA_HEADER not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_proxy_identity_as_a_signed_header( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + PROXY_IDENTITY, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'}, + ) + + assert json.loads(request.headers[REQUEST_METADATA_HEADER]) == EXPECTED_REQUEST_METADATA + assert REQUEST_METADATA_HEADER in _signed_header_names(request) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_keeps_caller_body_request_metadata_in_the_body( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**INVOKE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}} + + request: Final = await _sent_passthrough_request( + action, + caller_body, + PROXY_IDENTITY, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER.upper()}": '{"user_api_key_team_alias":"forged"}'}, + ) + + assert json.loads(request.content) == { + **INVOKE_BODY, + "requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"}, + } + assert REQUEST_METADATA_HEADER not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", [None, []]) +async def test_passthrough_requests_are_untouched_when_forwarding_is_off( + setting: list[str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", setting) + caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged"}} + + converse: Final = await _sent_passthrough_request("converse", caller_body, PROXY_IDENTITY) + invoke: Final = await _sent_passthrough_request("invoke", INVOKE_BODY, PROXY_IDENTITY) + + assert json.loads(converse.content) == caller_body + assert REQUEST_METADATA_HEADER not in invoke.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)]) +async def test_passthrough_drops_reserved_caller_keys_when_nothing_resolves( + action: str, body: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged"}} + + request: Final = await _sent_passthrough_request(action, caller_body, {}) + + assert json.loads(request.content) == body + assert REQUEST_METADATA_HEADER not in request.headers + + +MALFORMED_REQUEST_METADATA: Final = ( + pytest.param( + {"user_api_key_team_alias": "forged", "bad!": "x"}, "key 'bad!' must be", "ValidationException", id="bad-key" + ), + pytest.param({"ticket": "a\u00a0b"}, "value for key 'ticket' must be", "ValidationException", id="bad-value"), + pytest.param("ticket=T-42", "must be a dictionary", "SerializationException", id="not-a-dictionary"), + pytest.param({f"key_{index}": "v" for index in range(17)}, "maximum of 16 items", "ValidationException", id="17"), + pytest.param({"ticket": 42}, "must be strings", "SerializationException", id="non-string-value"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA) +async def test_malformed_converse_request_metadata_is_rejected_with_the_error_type_bedrock_returns( + action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match=violation) as rejected: + await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == converse_error_type + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA) +async def test_malformed_invoke_request_metadata_is_rejected_as_a_bedrock_validation_error( + action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match=violation) as rejected: + await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_null_request_metadata_is_rejected_like_bedrock_does( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match="must be a dictionary") as rejected: + await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": None}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +async def test_converse_null_request_metadata_is_treated_as_absent_like_bedrock_does( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": None}, PROXY_IDENTITY) + + assert json.loads(request.content) == {**CONVERSE_BODY, "requestMetadata": EXPECTED_REQUEST_METADATA} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_metadata_too_long_for_the_header_sends_the_leading_pairs_that_fit( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)} + resolved: Final = tuple( + {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform", **long_pairs}.items() + ) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + {**PROXY_IDENTITY, "spend_logs_metadata": long_pairs}, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'}, + ) + + header: Final = request.headers[REQUEST_METADATA_HEADER] + sent: Final = json.loads(header) + one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":")) + assert sent == dict(resolved[: len(sent)]) + assert len(header) <= 8500 < len(one_more_pair) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)]) +@pytest.mark.parametrize("drop_globally", [True, False]) +async def test_malformed_caller_request_metadata_is_dropped_when_params_are_dropped( + action: str, body: dict[str, object], drop_globally: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged", "bad!": "x"}} + + dropped: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True}) + ) + + assert json.loads(dropped.content) == {**body, "requestMetadata": EXPECTED_REQUEST_METADATA} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "body", "caller_metadata"), + [ + ("converse", CONVERSE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}), + ("invoke", INVOKE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}), + ("invoke", INVOKE_BODY, None), + ], +) +@pytest.mark.parametrize("drop_globally", [True, False]) +async def test_passthrough_body_is_sent_as_is_when_params_are_dropped_and_forwarding_is_off( + action: str, + body: dict[str, object], + caller_metadata: dict[str, str] | None, + drop_globally: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + caller_body: Final = {**body, "requestMetadata": caller_metadata} + + sent: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True}) + ) + + assert json.loads(sent.content) == caller_body + assert REQUEST_METADATA_HEADER not in sent.headers diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 8f95a502137..a181c5a5425 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -17,6 +17,7 @@ from litellm.constants import ( from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.realtime_api import main as realtime_main @pytest.fixture(autouse=True) @@ -109,6 +110,14 @@ class UnavailableBedrockStream: raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable") +class RecordingRealtimeHandler: + def __init__(self): + self.calls = 0 + + async def async_realtime(self, **kwargs): + self.calls += 1 + + class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id @@ -678,18 +687,23 @@ class TestBedrockRealtimeSessionLifecycle: @pytest.mark.asyncio @pytest.mark.parametrize( - "drop_params,expected_errors", - [(False, [("invalid_type", "session.temperature")]), (True, [])], + "session,drop_params,expected_errors", + [ + ({"temperature": "0.2", "modalities": ["text"]}, False, [("invalid_type", "session.temperature")]), + ({"temperature": "0.2", "modalities": ["text"]}, True, []), + ({"modalities": "text"}, False, [("invalid_type", "session.modalities")]), + ({"modalities": "text"}, True, []), + ], ) - async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers( - self, stub_aws_sdk_client, drop_params, expected_errors + async def test_malformed_setting_is_rejected_or_dropped_and_the_session_still_answers( + self, stub_aws_sdk_client, session, drop_params, expected_errors ): stream: Final = ScriptedBedrockStream([]) stub_aws_sdk_client["streams"] = [stream] user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} websocket: Final = ConnectedClientWS( [ - json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}), + json.dumps({"type": "session.update", "session": session}), json.dumps({"type": "conversation.item.create", "item": user_text}), ] ) @@ -777,12 +791,26 @@ class TestBedrockRealtimeProviderFailurePropagation: dispatched["logging_obj"] = RecordingLogging() return dispatched + USER_TEXT: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + } + ) + REJECTED_SESSION_UPDATE: Final = json.dumps({"type": "session.update", "session": {"temperature": "0.2"}}) + USER_AUDIO: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_messages", + [[SESSION_UPDATE], [USER_TEXT], [USER_AUDIO], [REJECTED_SESSION_UPDATE, USER_TEXT]], + ids=["session-update-first", "text-first", "audio-first", "input-after-rejected-update"], + ) async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( - self, stub_aws_sdk_client, spend_dispatch + self, stub_aws_sdk_client, spend_dispatch, monkeypatch, client_messages: list[str] ): handler = BedrockRealtime() - websocket = ConnectedClientWS([self.SESSION_UPDATE]) + websocket = ConnectedClientWS(client_messages) stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver) stub_aws_sdk_client["streams"] = [stream] @@ -814,6 +842,37 @@ class TestBedrockRealtimeProviderFailurePropagation: assert not litellm.should_retry(replay.value.status_code), "the router must not retry the replay refusal" assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error" + other_provider: Final = RecordingRealtimeHandler() + monkeypatch.setattr(realtime_main, "openai_realtime", other_provider) + with pytest.raises(BedrockError) as fallback: + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime", websocket=websocket, litellm_logging_obj=FakeLogging() + ) + + assert fallback.value.status_code == 400, "a fallback to another provider must not start a second session" + assert other_provider.calls == 0 + + @pytest.mark.asyncio + async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.USER_TEXT]) + healthy_stream = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + assert failure.value.status_code == 503 + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] + replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent] + assert any(event["event"].get("textInput", {}).get("content") == "hi" for event in replayed) + @pytest.mark.asyncio async def test_input_side_failure_keeps_spend_for_responses_already_delivered( self, stub_aws_sdk_client, spend_dispatch diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 1cc135a8671..3e10c3bcbdc 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -1655,9 +1655,17 @@ class TestBedrockRealtimeUserEventsAndUsage: ({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"), ({"max_response_output_tokens": True}, "session.max_response_output_tokens"), ({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"), + ({"instructions": {"role": "pirate"}}, "session.instructions"), + ({"voice": {"id": "voice_123"}}, "session.voice"), + ({"modalities": "text"}, "session.modalities"), + ({"modalities": ["text", 1]}, "session.modalities"), + ({"input_sample_rate_hertz": "16000"}, "session.input_sample_rate_hertz"), + ({"tools": {"type": "function", "name": "lookup"}}, "session.tools"), + ({"input_audio_format": {"type": "audio/pcm"}}, "session.input_audio_format"), + ({"output_audio_format": {"type": "audio/pcmu"}}, "session.output_audio_format"), ], ) - def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param): + def test_malformed_setting_is_rejected_without_drop_params(self, session, param): rejection: Final = BedrockRealtimeConfig().session_update_rejection( {"type": "session.update", "session": session} ) @@ -1677,6 +1685,59 @@ class TestBedrockRealtimeUserEventsAndUsage: inference: Final = self._inference_configuration(config, session) assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + def test_non_number_input_sample_rate_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"input_sample_rate_hertz": "8000"} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + config.transform_session_update_event({"type": "session.update", "session": session}) + assert config.input_sample_rate_hertz == BedrockRealtimeConfig().input_sample_rate_hertz + + def test_non_string_instructions_and_voice_are_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"instructions": {"role": "pirate"}, "voice": {"id": "voice_123"}} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + sent: Final = self._voice_and_system_prompt(config, session) + assert sent == self._voice_and_system_prompt(BedrockRealtimeConfig(), {}) + + @pytest.mark.parametrize("voice", ["alloy", "marin", "Alloy", "", " "]) + def test_openai_or_empty_voice_is_rejected_without_drop_params(self, voice): + rejection: Final = BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": {"voice": voice}} + ) + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_value", "session.voice") + assert f"'{voice}'" in rejection.message + + def test_openai_voice_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"voice": "alloy", "instructions": "Answer as a pirate."} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + voice, system_prompt = self._voice_and_system_prompt(config, session) + assert voice == self._voice_and_system_prompt(BedrockRealtimeConfig(), {})[0] + assert system_prompt == "Answer as a pirate." + + def test_nova_voice_is_applied(self): + session: Final = {"voice": "tiffany"} + assert BedrockRealtimeConfig().session_update_rejection({"type": "session.update", "session": session}) is None + assert self._voice_and_system_prompt(BedrockRealtimeConfig(), session)[0] == "tiffany" + + def test_g711_input_audio_format_sets_the_input_sample_rate(self): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"input_audio_format": "g711_ulaw"}}) + assert config.input_sample_rate_hertz == 8000 + + @staticmethod + def _voice_and_system_prompt(config: BedrockRealtimeConfig, session: dict[str, object]) -> tuple[object, object]: + events: Final = [ + json.loads(message)["event"] + for message in config.transform_session_update_event({"type": "session.update", "session": session}) + ] + voice: Final = next( + e["promptStart"]["audioOutputConfiguration"]["voiceId"] for e in events if "promptStart" in e + ) + system_prompt: Final = next(e["textInput"]["content"] for e in events if "textInput" in e) + return voice, system_prompt + def test_null_number_settings_count_as_unset(self): session: Final = {"temperature": None, "max_response_output_tokens": None} config: Final = BedrockRealtimeConfig() diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 303f752ab0d..9aab5052cfe 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -659,8 +659,14 @@ def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_m apply_bedrock_invoke_structured_output, ) - legacy = {"type": "json_schema", "schema": {"type": "object", "properties": {"a": {"type": "string"}}}} - newer = {"type": "json_schema", "schema": {"type": "object", "properties": {"b": {"type": "string"}}}} + legacy = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"a": {"type": "string"}}, "additionalProperties": False}, + } + newer = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"b": {"type": "string"}}, "additionalProperties": False}, + } body = { "messages": [{"role": "user", "content": "hi"}], "output_format": legacy, diff --git a/tests/unit/llms/bedrock/test_request_metadata.py b/tests/unit/llms/bedrock/test_request_metadata.py index 5a14bbea9f2..6b74ad90cff 100644 --- a/tests/unit/llms/bedrock/test_request_metadata.py +++ b/tests/unit/llms/bedrock/test_request_metadata.py @@ -1,8 +1,10 @@ import asyncio import json +from collections.abc import Callable, Mapping +from typing import Final import pytest - +from pydantic import TypeAdapter import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -18,6 +20,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran ) from litellm.llms.bedrock.request_metadata import ( BEDROCK_REQUEST_METADATA_HEADER, + BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS, BEDROCK_REQUEST_METADATA_MAX_PAIRS, resolve_bedrock_request_metadata, ) @@ -128,6 +131,20 @@ def test_identity_survives_a_caller_filling_every_slot( ) +def test_a_key_sent_twice_keeps_the_caller_value_and_uses_one_slot(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) + caller: Final = { + f"client_{index:02d}": "caller" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY) - 1) + } + spend_logs: Final = {**{key: "spend" for key in caller}, "extra": "spend"} + + resolved: Final = resolve_bedrock_request_metadata( + litellm_params("metadata", spend_logs_metadata=spend_logs, **IDENTITY), caller_metadata=caller + ) + + assert resolved == {**IDENTITY, **caller, "extra": "spend"} + + @pytest.mark.parametrize( "field_order", [ @@ -158,7 +175,14 @@ def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot( @pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"]) @pytest.mark.parametrize( "forged_key", - ["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"], + [ + "user_api_key_team_alias", + "user_api_key_org_alias", + "user_api_key_hash", + "USER_API_KEY_TEAM_ALIAS", + "User_Api_Key_Hash", + " user_api_key_user_id", + ], ) def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( forged_key, @@ -181,14 +205,21 @@ def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( assert "attacker-controlled" not in resolved.values() -def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize( + "team_alias", + ["O'Brien's team", "platform\u00a0team", "platform\u3000team", "x" * 256 + "\n"], + ids=["apostrophe", "no-break-space", "ideographic-space", "257-chars-ending-in-newline"], +) +def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds( + team_alias: str, monkeypatch: pytest.MonkeyPatch +) -> None: """A team alias with an apostrophe must not turn a working request into a 400 the moment - an operator flips the setting on.""" + an operator flips the setting on. Bedrock also rejects non-ASCII whitespace (checked live 2026-10-16).""" monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) params = litellm_params( "metadata", user_api_key_alias="prod-key", - user_api_key_team_alias="O'Brien's team", + user_api_key_team_alias=team_alias, user_api_key_user_email="x" * 300, ) @@ -198,13 +229,54 @@ def test_identity_violating_the_character_class_is_dropped_and_the_request_succe assert body["messages"] -def test_caller_supplied_violation_still_raises_bad_request(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize( + "caller_metadata", + [{"team": "O'Brien's team"}, {"k" * 256 + "\n": "v"}, {"team": "v" * 256 + "\n"}], + ids=["bad-character", "257-char-key-ending-in-newline", "257-char-value-ending-in-newline"], +) +def test_caller_supplied_violation_still_raises_bad_request( + caller_metadata: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) with pytest.raises(litellm.exceptions.BadRequestError): - converse_body( - litellm_params("metadata", **IDENTITY), - {"requestMetadata": {"team": "O'Brien's team"}}, + converse_body(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_metadata}) + + +@pytest.mark.parametrize("driver", CONVERSE_DRIVERS) +@pytest.mark.parametrize("drop_globally", [True, False]) +@pytest.mark.parametrize(("fields", "forwarded_identity"), [(ALL_FIELDS, IDENTITY), (None, {})]) +def test_caller_violations_are_dropped_like_the_passthrough_routes_when_params_are_dropped( + driver: Callable[[Mapping[str, object], Mapping[str, object]], Mapping[str, object]], + drop_globally: bool, + fields: list[str] | None, + forwarded_identity: Mapping[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + params: Final = {**litellm_params("metadata", **IDENTITY), **({} if drop_globally else {"drop_params": True})} + + body: Final = driver(params, {"requestMetadata": {"team": "O'Brien's team", "ticket": "T-42"}}) + + assert body["requestMetadata"] == {**forwarded_identity, "ticket": "T-42"} + + +@pytest.mark.parametrize("fields", [ALL_FIELDS, None]) +def test_mapped_caller_violations_are_dropped_when_params_are_dropped( + fields: list[str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields) + caller: Final = {"team": "O'Brien's team", "ticket": "T-42"} + + mapped: Final = AmazonConverseConfig().map_openai_params( + non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=True + ) + + assert mapped["requestMetadata"] == caller + with pytest.raises(litellm.exceptions.BadRequestError): + AmazonConverseConfig().map_openai_params( + non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=False ) @@ -320,6 +392,27 @@ def metadata_header_values(headers): return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()] +@pytest.mark.parametrize("driver", HEADER_DRIVERS) +def test_the_metadata_header_keeps_the_leading_pairs_that_fit_bedrocks_header_limit( + driver: Callable[[Mapping[str, str], Mapping[str, object]], Mapping[str, str]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) + long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)} + params: Final = litellm_params("metadata", spend_logs_metadata=long_pairs, **IDENTITY) + resolved: Final = tuple((resolve_bedrock_request_metadata(params) or {}).items()) + + headers: Final = driver({}, params) + (header,) = (value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()) + sent: Final = TypeAdapter(dict[str, str]).validate_json(header) + one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":")) + + assert len(header) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + assert tuple(sent.items()) == resolved[: len(sent)] + assert IDENTITY.items() <= sent.items() + assert len(one_more_pair) > BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + + def test_converse_still_sets_the_bearer_authorization_header(monkeypatch: pytest.MonkeyPatch): """Converse owns the metadata header now, and that must not disturb the api_key path its validate_environment existed for. Closing the forgery hole cannot break authentication.""" diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 50ed70e09f5..23c142b61ab 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5313,3 +5313,45 @@ async def test_retrieve_file_rejects_unverified_bedrock_range_error( ) finally: sync_client.close() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_its_beta() -> None: + received: Final[list[httpx.Request]] = [] + + def anthropic_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "msg_dropped_safeguards", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(anthropic_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="anthropic/claude-sonnet-4-6", + api_key="sk-ant-test", + safeguards=({"type": "dangerous_tool_use"},), + additional_drop_params=["safeguards"], + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + sent: Final = received[0] + assert "safeguards" not in json.loads(sent.content) + assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) diff --git a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 5cbb9a0a7b4..46062f16c9a 100644 --- a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock import httpx @@ -9,6 +10,10 @@ from litellm.llms.github_copilot.common_utils import GetAPIKeyError from litellm.llms.github_copilot.messages.transformation import ( GithubCopilotAnthropicMessagesConfig, ) +from litellm.llms.github_copilot.per_user_auth import ( + GITHUB_COPILOT_USER_SESSION_KWARG_KEY, + GithubCopilotUserSession, +) def test_github_copilot_anthropic_messages_get_complete_url(): @@ -185,6 +190,26 @@ def test_github_copilot_anthropic_messages_validate_environment_injects_context_ assert "context-management-2025-06-27" in validated_headers["anthropic-beta"] +@pytest.mark.usefixtures("local_beta_headers_config") +def test_github_copilot_safeguards_do_not_derive_the_dangerous_tool_use_beta() -> None: + session: Final = GithubCopilotUserSession(token="gh.test-key", api_base="https://api.githubcopilot.com") + + validated_headers, _ = GithubCopilotAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="github_copilot/claude-haiku-4.5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={ + "safeguards": ({"type": "dangerous_tool_use"},), + "context_management": {"edits": ({"type": "clear_tool_uses_20250919"},)}, + }, + litellm_params={GITHUB_COPILOT_USER_SESSION_KWARG_KEY: session}, + api_key=None, + api_base=None, + ) + + assert validated_headers["anthropic-beta"] == "context-management-2025-06-27" + + def test_github_copilot_anthropic_messages_validate_environment_auth_error(): """Test error handling when authentication fails.""" config = GithubCopilotAnthropicMessagesConfig() diff --git a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py index 86fb6fa15d8..9b8788d3769 100644 --- a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py +++ b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py @@ -1,15 +1,16 @@ import json +import re from io import BytesIO from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest from openai import AsyncOpenAI, OpenAI +from openai._legacy_response import LegacyAPIResponse from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage -import pytest -import re import litellm from litellm import completion @@ -18,7 +19,6 @@ from litellm.llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfi from tests._vcr_conftest_common import rewound_new_episodes_cassette from tests.capturing_transport import CapturingTransport - _GATEWAY_EMBEDDING_RESPONSE: Final = CreateEmbeddingResponse( object="list", data=(Embedding(object="embedding", index=0, embedding=(0.1, 0.2, 0.3)),), @@ -98,7 +98,7 @@ async def test_litellm_gateway_from_sdk(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions.with_raw_response, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -114,9 +114,9 @@ async def test_litellm_gateway_from_sdk(): mock_call.assert_called_once() - print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) + print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"])) - assert "hello" in mock_call.call_args.kwargs["extra_body"] + assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"] @pytest.mark.parametrize("is_async", (False, True)) @@ -264,7 +264,7 @@ async def test_litellm_gateway_from_sdk_rerank(is_async): patch_target = client.post with patch.object(client, "post", new=mock_method): - mock_response = MagicMock() + mock_response = MagicMock(spec=LegacyAPIResponse) # Create a mock response similar to OpenAI's rerank response mock_response.text = json.dumps( @@ -408,7 +408,7 @@ async def test_litellm_gateway_from_sdk_structured_output(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: litellm.completion( @@ -427,8 +427,8 @@ async def test_litellm_gateway_from_sdk_structured_output(): mock_call.assert_called_once() - print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) - json_schema = mock_call.call_args.kwargs["response_format"] + print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"])) + json_schema = mock_call.call_args.kwargs["body"]["response_format"] assert "json_schema" in json_schema @@ -486,7 +486,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): openai_client = OpenAI(api_key="fake-key") # Create mock response object - mock_response = MagicMock() + mock_response = MagicMock(spec=LegacyAPIResponse) mock_response.headers = {"x-litellm-response-cost": "120"} mock_response.parse.return_value = litellm.ModelResponse( **{ @@ -528,8 +528,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): ) with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", return_value=mock_response, ) as mock_call: response = litellm.completion( @@ -564,7 +563,7 @@ async def test_litellm_gateway_image_generation_direct(is_async): "created": 1, "data": [{"url": "https://example.com/image.png"}], } - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.parse.return_value = mock_openai_response mock_raw_response.headers = {} diff --git a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py index e5298c7ba39..3633ac97265 100644 --- a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py +++ b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py @@ -94,7 +94,7 @@ def test_chat_completion_nvidia_nim_with_tools(): ] with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: completion( @@ -116,7 +116,7 @@ def test_chat_completion_nvidia_nim_with_tools(): # Add assertions to check the request mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) @@ -208,7 +208,7 @@ def test_completion_nvidia_nim(): api_key="fake-api-key", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: completion( model=model_name, @@ -226,7 +226,7 @@ def test_completion_nvidia_nim(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["messages"] == [ { diff --git a/tests/unit/llms/openai/test_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py index 6ee882b7d32..121fab66c15 100644 --- a/tests/unit/llms/openai/test_o_series_transformation.py +++ b/tests/unit/llms/openai/test_o_series_transformation.py @@ -70,7 +70,7 @@ async def test_o1_handle_system_role(model): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -83,7 +83,7 @@ async def test_o1_handle_system_role(model): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) @@ -141,7 +141,7 @@ async def test_o1_max_completion_tokens(model: str): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -154,7 +154,7 @@ async def test_o1_max_completion_tokens(model: str): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index d7ca9319f3a..11d7c7d6d3c 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -5,12 +5,13 @@ import os from collections.abc import Mapping from types import ModuleType from typing import Final, Literal, Protocol -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import httpx import pytest import respx from openai import AsyncOpenAI, AzureOpenAI, OpenAI +from openai._legacy_response import LegacyAPIResponse from openai._types import Response as SDKResponse from pydantic import BaseModel, ConfigDict, JsonValue @@ -1669,7 +1670,7 @@ async def test_openai_prediction_param_mock(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1688,7 +1689,7 @@ async def test_openai_prediction_param_mock(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the prediction parameter assert "prediction" in request_body @@ -1700,10 +1701,12 @@ async def test_openai_prediction_param_mock(): def test_openai_max_retries_0(mock_get_openai_client): import litellm - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( - ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) - ) + raw_response: Final = MagicMock(spec=LegacyAPIResponse) + raw_response.headers = {} + raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + client_stub: Final = MagicMock() + client_stub.post.return_value = raw_response + mock_get_openai_client.return_value = client_stub litellm.set_verbose = True response = litellm.completion( model="gpt-4o-mini", @@ -1812,7 +1815,7 @@ async def test_openai_safety_identifier_parameter(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1825,7 +1828,7 @@ async def test_openai_safety_identifier_parameter(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the safety_identifier parameter assert "safety_identifier" in request_body @@ -1841,7 +1844,7 @@ def test_openai_safety_identifier_parameter_sync(): client = OpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: litellm.completion( @@ -1854,7 +1857,7 @@ def test_openai_safety_identifier_parameter_sync(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the safety_identifier parameter assert "safety_identifier" in request_body @@ -1871,7 +1874,7 @@ async def test_openai_service_tier_parameter(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1884,7 +1887,7 @@ async def test_openai_service_tier_parameter(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the service_tier parameter assert "service_tier" in request_body, "service_tier should be in request body" @@ -1902,7 +1905,7 @@ def test_openai_service_tier_parameter_sync(): client = OpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: litellm.completion( @@ -1915,7 +1918,7 @@ def test_openai_service_tier_parameter_sync(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the service_tier parameter assert "service_tier" in request_body, "service_tier should be in request body" @@ -1947,3 +1950,63 @@ def test_responses_gpt54_with_xhigh_reasoning(): assert request_body["model"] == "openai/gpt-5.4" assert request_body["reasoning"] == {"effort": "xhigh"} + + +_CHAT_RESPONSE_JSON: Final = { + "id": "chatcmpl-wire", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_WIRE_REQUEST: Final = { + "model": "gpt-5.4", + "messages": [ + {"role": "system", "content": "sys"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is in the image"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png", "detail": "low"}}, + ], + }, + ], + "tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}], + "tool_choice": "auto", + "response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}, + "max_completion_tokens": 7, + "reasoning_effort": "none", + "n": 2, + "seed": 3, +} + + +@pytest.mark.parametrize("is_async", (False, True)) +@respx.mock +async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_headers_unchanged( + is_async: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_CHAT_RESPONSE_JSON) + ) + kwargs: Final = { + **_WIRE_REQUEST, + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1", + "api_key": "sk-test", + "num_retries": 0, + "extra_body": {"vendor_flag": True}, + "extra_headers": {"x-caller": "unit"}, + } + + if is_async: + await litellm.acompletion(**kwargs) + else: + litellm.completion(**kwargs) + + sent: Final = route.calls.last.request + assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True} + assert sent.headers["x-caller"] == "unit" + assert sent.headers["authorization"] == "Bearer sk-test" diff --git a/tests/unit/llms/openai/test_openai_empty_response.py b/tests/unit/llms/openai/test_openai_empty_response.py index 26b28f967db..1505d9ef049 100644 --- a/tests/unit/llms/openai/test_openai_empty_response.py +++ b/tests/unit/llms/openai/test_openai_empty_response.py @@ -2,13 +2,13 @@ Test for issue #17209: Clearer error when LLM endpoint returns empty response """ -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest +from openai._legacy_response import LegacyAPIResponse - -from litellm.llms.openai.openai import OpenAIChatCompletion from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.openai import OpenAIChatCompletion class TestEmptyResponseHandling: @@ -22,12 +22,12 @@ class TestEmptyResponseHandling: openai_chat = OpenAIChatCompletion() # Mock the raw response to return an empty string from parse() - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = "" # Empty string response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -48,12 +48,12 @@ class TestEmptyResponseHandling: """Test that None response also produces a clear error""" openai_chat = OpenAIChatCompletion() - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = None mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -75,12 +75,12 @@ class TestEmptyResponseHandling: mock_response = MagicMock() mock_response.model_dump.return_value = {"choices": []} - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {"x-request-id": "123"} mock_raw_response.parse.return_value = mock_response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -104,12 +104,12 @@ class TestEmptyResponseHandling: # Create a mock response WITHOUT model_dump (like an AsyncStream/Iterator) mock_stream = MagicMock(spec=[]) # spec=[] means no attributes - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {"x-request-id": "123"} mock_raw_response.parse.return_value = mock_stream mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) diff --git a/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py index 7013afc7a5f..58f1ff60ba5 100644 --- a/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py +++ b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py @@ -9,6 +9,7 @@ body. OpenAI/Anthropic reject unknown body params with HTTP 400. from unittest.mock import MagicMock +from openai._legacy_response import LegacyAPIResponse import litellm from litellm.types.utils import all_litellm_params @@ -47,12 +48,12 @@ def test_completion_does_not_leak_flag_into_provider_request_body(): }, } - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = mock_response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -65,7 +66,7 @@ def test_completion_does_not_leak_flag_into_provider_request_body(): ) create_kwargs = ( - mock_client.chat.completions.with_raw_response.create.call_args.kwargs + mock_client.post.call_args.kwargs ) assert "use_chat_completions_api" not in create_kwargs assert "use_chat_completions_api" not in (create_kwargs.get("extra_body") or {}) diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py index 88382b80d16..dbc6ec54a40 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py @@ -1,5 +1,11 @@ +from typing import Final +from unittest.mock import MagicMock + +import httpx import pytest + from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.types.utils import ModelResponse def test_sanitize_usage_obj_handles_null_tokens(): @@ -47,3 +53,32 @@ def test_sanitize_usage_obj_valid_usage(): # Assert assert sanitized_json == original_json # The object should be unchanged + + +def test_transform_response_surfaces_provider_headers_as_llm_provider_headers(): + raw_response: Final = httpx.Response( + status_code=200, + headers={"x-amzn-requestid": "req_abc123", "content-type": "application/json"}, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "google.gemma-4-31b", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + response: Final = OpenAILikeChatConfig().transform_response( + model="google.gemma-4-31b", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req_abc123" diff --git a/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 9167853bf64..1d5ef50ca3a 100644 --- a/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -516,3 +516,26 @@ def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(co assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} + + +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("client_headers", "expected_beta"), + [ + ({}, None), + ({"anthropic-beta": "dangerous-tool-use-2026-09-03"}, "dangerous-tool-use-2026-09-03"), + ], +) +def test_safeguards_reach_an_openai_like_host_with_only_the_client_betas( + config: OpenAILikeAnthropicMessagesConfig, client_headers: dict[str, str], expected_beta: str | None +) -> None: + headers, _ = config.validate_anthropic_messages_environment( + headers=client_headers, + model="some-model", + messages=[{"role": "user", "content": "hi"}], + optional_params={"safeguards": ({"type": "dangerous_tool_use"},)}, + litellm_params={}, + api_key="sk-test", + api_base="https://host/v1", + ) + assert headers.get("anthropic-beta") == expected_beta, headers diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 82ad7702052..0dba0d7b90f 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -1,7 +1,8 @@ import json -from typing import Final +from typing import Final, Literal import pytest +from typing_extensions import ReadOnly, TypedDict from litellm.constants import SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler @@ -50,3 +51,162 @@ def test_full_payload_guardrail_text_scans_caller_value_but_not_the_marker(marke ) assert SERVER_STREAMING_CLASSIFICATION_KEY not in json.loads(marker_text), marker_text + + +def _converse_frame(event_type: str, payload: dict) -> bytes: + import struct + from binascii import crc32 + + headers = b"".join( + bytes((len(name),)) + name + b"\x07" + struct.pack("!H", len(value)) + value + for name, value in ( + (b":event-type", event_type.encode()), + (b":content-type", b"application/json"), + (b":message-type", b"event"), + ) + ) + body = json.dumps(payload).encode() + prelude = struct.pack("!II", 16 + len(headers) + len(body), len(headers)) + message = prelude + struct.pack("!I", crc32(prelude)) + headers + body + return message + struct.pack("!I", crc32(message)) + + +async def _frames(*raw: bytes) -> list[object]: + from litellm.llms.bedrock.passthrough.guardrail_translation.handler import BedrockPassthroughGuardrailHandler + + async def chunks(): + for chunk in raw: + yield chunk + + return [frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(chunks(), "converse-stream")] + + +@pytest.mark.asyncio +async def test_passthrough_route_scans_converse_frames_with_the_bedrock_translation(): + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + frames: Final = await _frames( + _converse_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "Hello "}}), + _converse_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "world"}}), + _converse_frame("messageStop", {"stopReason": "end_turn"}), + ) + + scan_key: Final = LlmPassthroughRouteHandler().get_streaming_scan_key(frames) + + assert scan_key is not None and (scan_key.texts, scan_key.stream_ended) == (("Hello world",), True) + assert LlmPassthroughRouteHandler()._check_streaming_has_ended(frames) is True + assert LlmPassthroughRouteHandler()._check_streaming_has_ended(frames[:2]) is False + + +def test_passthrough_route_has_no_stream_translation_for_items_no_provider_owns(): + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + assert LlmPassthroughRouteHandler().get_streaming_scan_key([b"raw", {"chunk": 1}]) is None + assert LlmPassthroughRouteHandler()._check_streaming_has_ended([b"raw"]) is False + + +@pytest.mark.parametrize(("provider", "ends_message"), [("bedrock", True), ("vertex_ai", False)]) +def test_passthrough_route_block_before_any_frame_uses_the_request_provider(provider: str, ends_message: bool): + from litellm.exceptions import ModifyResponseException + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + chunks: Final = LlmPassthroughRouteHandler().build_block_sse_chunks( + ModifyResponseException(message="blocked", model="m", request_data={"custom_llm_provider": provider}), + stream_started=False, + responses_so_far=[], + ) + + assert (chunks is not None) is ends_message + + +_INVOKE_STREAM: Final = "model/anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream" + + +class _TextDelta(TypedDict): + type: ReadOnly[Literal["text_delta"]] + text: ReadOnly[str] + + +class _TextDeltaEvent(TypedDict): + type: ReadOnly[Literal["content_block_delta"]] + index: ReadOnly[int] + delta: ReadOnly[_TextDelta] + + +def _invoke_frame(event: _TextDeltaEvent) -> bytes: + import base64 + + return _converse_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +@pytest.mark.parametrize( + ("provider", "endpoint", "model", "body", "relayable"), + [ + ("bedrock", "model/x/converse-stream", "meta.llama4-maverick-17b-instruct-v1:0", None, True), + ("bedrock", _INVOKE_STREAM, "anthropic.claude-sonnet-5-v1:0", None, True), + ("bedrock", _INVOKE_STREAM, "meta.llama4-maverick-17b-instruct-v1:0", {"prompt": "hi"}, False), + ("bedrock", _INVOKE_STREAM, None, None, False), + ("bedrock", _INVOKE_STREAM, None, {"anthropic_version": "bedrock-2023-05-31"}, True), + ("bedrock", None, "anthropic.claude-sonnet-5-v1:0", None, False), + ("vertex_ai", "model/x/converse-stream", "anthropic.claude-sonnet-5-v1:0", None, False), + ], +) +def test_passthrough_route_relays_and_buffers_the_streams_its_provider_translation_scans( + provider: str, endpoint: str | None, model: str | None, body: object, relayable: bool +): + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + assert LlmPassthroughRouteHandler.supports_event_stream_relay(provider, endpoint, model, body) is relayable + assert ( + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization(provider, endpoint, model, body) is relayable + ) + + +@pytest.mark.asyncio +async def test_passthrough_route_frames_an_invoke_stream_for_the_bedrock_translation(): + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + async def chunks(): + yield _invoke_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi "}}) + yield _invoke_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "you"}}) + + items: Final = LlmPassthroughRouteHandler.event_stream_items("bedrock", _INVOKE_STREAM, chunks()) + assert items is not None + frames: Final = [frame async for frame in items] + + scan_key: Final = LlmPassthroughRouteHandler().get_streaming_scan_key(frames) + + assert scan_key is not None and scan_key.texts == ("Hi you",) + assert LlmPassthroughRouteHandler()._check_streaming_has_ended(frames) is True + assert LlmPassthroughRouteHandler()._check_streaming_has_ended(frames[:2]) is False + + +def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_chunk_frames(): + import base64 + + from botocore.eventstream import EventStreamBuffer + + from litellm.exceptions import ModifyResponseException + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + chunks: Final = LlmPassthroughRouteHandler().build_block_sse_chunks( + ModifyResponseException( + message="blocked", model="m", request_data={"custom_llm_provider": "bedrock", "endpoint": _INVOKE_STREAM} + ), + stream_started=False, + responses_so_far=[], + ) + + assert chunks is not None + buffer: Final = EventStreamBuffer() + buffer.add_data(b"".join(chunks)) + events: Final = [json.loads(base64.b64decode(json.loads(message.payload)["bytes"])) for message in buffer] + assert [event["type"] for event in events] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert events[2]["delta"]["text"] == "blocked" diff --git a/tests/unit/llms/perplexity/chat/test_transformation.py b/tests/unit/llms/perplexity/chat/test_transformation.py index 6a70fd1232d..0cb6335f5a8 100644 --- a/tests/unit/llms/perplexity/chat/test_transformation.py +++ b/tests/unit/llms/perplexity/chat/test_transformation.py @@ -4,6 +4,7 @@ import os from unittest.mock import MagicMock, patch import pytest +from openai._legacy_response import LegacyAPIResponse import litellm from litellm import completion @@ -92,7 +93,7 @@ class TestPerplexityReasoning: pydantic_obj = ChatCompletion(**response_object) def _return_pydantic_obj(*args, **kwargs): - new_response = MagicMock() + new_response = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"content-type": "application/json"} new_response.parse.return_value = pydantic_obj return new_response @@ -100,8 +101,7 @@ class TestPerplexityReasoning: openai_client = OpenAI(api_key="fake-api-key") with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", side_effect=_return_pydantic_obj, ) as mock_client: response = completion( @@ -121,7 +121,7 @@ class TestPerplexityReasoning: # Get the request data from the mock call call_args = mock_client.call_args - request_data = call_args.kwargs + request_data = call_args.kwargs["body"] # Verify reasoning_effort was included in the request assert "reasoning_effort" in request_data diff --git a/tests/unit/llms/volcengine/test_volcengine.py b/tests/unit/llms/volcengine/test_volcengine.py index d472d52d2db..dff9a0c62b2 100644 --- a/tests/unit/llms/volcengine/test_volcengine.py +++ b/tests/unit/llms/volcengine/test_volcengine.py @@ -2,6 +2,7 @@ import os import sys from unittest.mock import MagicMock, patch +from openai._legacy_response import LegacyAPIResponse from pydantic import BaseModel from litellm.llms.volcengine.chat.transformation import ( @@ -117,7 +118,7 @@ class TestVolcEngineConfig: client = OpenAI(api_key="test_api_key") - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = { "x-request-id": "123", "openai-organization": "org-123", @@ -126,9 +127,7 @@ class TestVolcEngineConfig: } mock_raw_response.parse.return_value = ModelResponse() - with patch.object( - client.chat.completions.with_raw_response, "create", mock_raw_response - ) as mock_create: + with patch.object(client, "post", MagicMock(return_value=mock_raw_response)) as mock_create: completion( model="volcengine/doubao-seed-1.6", messages=[ @@ -144,11 +143,6 @@ class TestVolcEngineConfig: ) mock_create.assert_called_once() - print(mock_create.call_args.kwargs) + print(mock_create.call_args.kwargs["body"]) # Fixed: thinking disabled should appear in extra_body with original structure - assert ( - "extra_body" in mock_create.call_args.kwargs - and "thinking" in mock_create.call_args.kwargs.get("extra_body", {}) - and mock_create.call_args.kwargs.get("extra_body", {})["thinking"] - == {"type": "disabled"} - ) + assert mock_create.call_args.kwargs["options"]["extra_json"]["thinking"] == {"type": "disabled"} diff --git a/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py index 922643f9834..4819688f14e 100644 --- a/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py @@ -128,6 +128,119 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == chunks[0] +@pytest.mark.asyncio +async def test_asyncpassthroughstreamingresponse_wait_for_spend_flush_returns_once_the_flush_has_finished(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + mock_response = _make_streaming_response([b"chunk-1", b"chunk-2"]) + + async def response_coro(): + return mock_response + + flushed: List[bytes] = [] + + async def flush(collector: _RecordingCollector) -> bool: + await asyncio.sleep(0) + await asyncio.sleep(0) + flushed.append(collector.build_logged_response(litellm_logging_obj=MagicMock())) + return True + + mock_logging_obj = MagicMock() + mock_logging_obj.async_flush_passthrough_collected_chunks = flush + gen = AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=_provider_config(_RecordingCollector()), + ) + + await gen.__anext__() + await gen.aclose() + + assert await gen.wait_for_spend_flush() is True + assert flushed == [b"chunk-1"] + + +@pytest.mark.asyncio +async def test_asyncpassthroughstreamingresponse_says_it_logged_no_spend_when_the_flush_had_nothing_to_log(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + async def response_coro(): + return _make_streaming_response([b"part-of-a-frame", b"rest-of-the-frame"]) + + mock_logging_obj = MagicMock() + mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock(return_value=False) + gen = AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=_provider_config(_RecordingCollector()), + ) + + await gen.__anext__() + await gen.aclose() + + assert await gen.wait_for_spend_flush() is False + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_asyncpassthroughstreamingresponse_closed_before_any_chunk_says_it_logged_no_spend(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + async def response_coro(): + return _make_streaming_response([b"chunk-1"]) + + mock_logging_obj = _make_logging_obj() + gen = AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=_provider_config(_RecordingCollector()), + ) + + await gen + await gen.aclose() + + assert await gen.wait_for_spend_flush() is False + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called() + + +@pytest.mark.asyncio +async def test_asyncpassthroughstreamingresponse_spend_flush_survives_a_cancelled_wait(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + mock_response = _make_streaming_response([b"chunk-1", b"chunk-2"]) + + async def response_coro(): + return mock_response + + release = asyncio.Event() + flushed: List[bytes] = [] + + async def flush(collector: _RecordingCollector) -> None: + await release.wait() + flushed.append(collector.build_logged_response(litellm_logging_obj=MagicMock())) + + mock_logging_obj = MagicMock() + mock_logging_obj.async_flush_passthrough_collected_chunks = flush + gen = AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=_provider_config(_RecordingCollector()), + ) + + await gen.__anext__() + await gen.aclose() + waiter = asyncio.create_task(gen.wait_for_spend_flush()) + await asyncio.sleep(0) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + release.set() + await gen.wait_for_spend_flush() + + assert flushed == [b"chunk-1"] + + @pytest.mark.asyncio async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx(): from litellm.passthrough.main import AsyncPassthroughStreamingResponse diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py index bb228544e14..ab2b5b3d966 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py @@ -13,6 +13,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion refresh_token_exchange_targets, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import SSOClientConfig from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SAML2_SUBJECT_TOKEN_TYPE, SSOIdentityAssertion, @@ -163,6 +164,90 @@ async def test_an_oidc_login_is_exchanged_at_the_id_jag_client_and_keeps_its_sso assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" +def _sso_client(client_id: str = "agent", endpoint: str = "https://idp.example.com/token") -> SSOClientConfig: + return SSOClientConfig( + token_endpoint=endpoint, client_id=client_id, client_secret=SecretStr("s"), auth_method="client_secret_basic" + ) + + +@pytest.mark.asyncio +async def test_an_oidc_login_through_the_id_jag_client_keeps_its_own_refresh_token_without_an_exchange(): + """Okta refuses requested_token_type=refresh_token for an id_token subject. When the SSO app is the + ID-JAG client itself, the refresh token the login returned is already that client's, so it is bound + to the target and nothing is posted to the IdP.""" + endpoint = _Endpoint( + Error( + CredError.of_upstream_unavailable( + "token exchange failed with status 400: 'requested_token_type' is invalid or not supported." + ) + ) + ) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + kept = await exchange_login_assertion_for_refresh_token( + oidc, (_config(),), endpoint, _NOW, sso_client=_sso_client() + ) + + assert endpoint.calls == [] + assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt" + assert kept.exchanged_refresh_token_expires_at is None + assert kept.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert kept.exchanged_refresh_token_client_id == "agent" + assert kept.id_token == oidc.id_token and kept.refresh_token == oidc.refresh_token + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sso_client", + [ + pytest.param(_sso_client(client_id="sso-app"), id="other-client"), + pytest.param(_sso_client(endpoint="https://other-idp.example.com/token"), id="other-endpoint"), + pytest.param(None, id="unknown-client"), + ], +) +async def test_a_login_refresh_token_from_another_client_or_issuer_is_never_presented( + sso_client: SSOClientConfig | None, +): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + exchanged = await exchange_login_assertion_for_refresh_token( + oidc, (_config(),), endpoint, _NOW, sso_client=sso_client + ) + + assert [(call[0], call[1], call[2]["subject_token"]) for call in endpoint.calls] == [ + ("https://idp.example.com/token", "agent", "id-token") + ] + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + +@pytest.mark.asyncio +async def test_the_login_refresh_token_binds_to_the_target_whose_client_issued_it_not_the_first(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + targets = (_config(client_id="other-agent"), _config(client_id="agent")) + + kept = await exchange_login_assertion_for_refresh_token(oidc, targets, endpoint, _NOW, sso_client=_sso_client()) + + assert endpoint.calls == [] + assert kept.exchanged_refresh_token_client_id == "agent" + assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt" + + +@pytest.mark.asyncio +async def test_a_saml_login_through_the_id_jag_client_still_exchanges_its_assertion(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + + exchanged = await exchange_login_assertion_for_refresh_token( + _saml(), (_config(),), endpoint, _NOW, sso_client=_sso_client() + ) + + assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:saml2" + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + @pytest.mark.asyncio async def test_a_login_already_exchanged_is_not_exchanged_again(): endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-2", expires_in=3600))) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index af62594e3b8..3c8e7a9b448 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -612,6 +612,27 @@ async def test_login_refresh_token_round_trips_beside_the_saml_assertion(): assert login_refresh_token_expired(fetched, now + timedelta(seconds=3600)) +def test_renewal_follows_a_rotated_login_refresh_token_that_is_the_id_jag_subject(): + """When the OIDC login's own refresh token is the ID-JAG subject (the SSO client is the ID-JAG + client), a renewal that rotated it leaves the old token revoked at the IdP, so the subject moves + to the token the renewal returned.""" + now = datetime.now(timezone.utc).replace(microsecond=0) + login = assertion_from_sso_login(_make_id_token(), "sso-rt") + assert login is not None + previous = with_login_refresh_token( + login, "sso-rt", None, now, endpoint="https://idp.example.com/token", client_id="agent" + ) + renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2") + assert renewed is not None + + carried = carry_exchanged_refresh_token(previous, renewed) + + assert carried.exchanged_refresh_token is not None + assert carried.exchanged_refresh_token.get_secret_value() == "sso-rt-2" + assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert carried.exchanged_refresh_token_client_id == "agent" + + def test_renewal_carries_the_exchanged_refresh_token_forward(): """An OIDC renewal rebuilds the row from the new id_token. The refresh token the ID-JAG client exchanged at login is unrelated to that renewal and must survive it.""" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 5fdc7669496..a43dacaf100 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -5,11 +5,12 @@ import json import time from base64 import urlsafe_b64encode from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException, Request +from fastapi import HTTPException, Request, Response +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.types.mcp import MCPAuth @@ -26,6 +27,14 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +class _BridgeRegistrationPayload(TypedDict, total=False): + client_name: ReadOnly[str] + grant_types: ReadOnly[list[str]] + response_types: ReadOnly[list[str]] + token_endpoint_auth_method: ReadOnly[str] + redirect_uris: ReadOnly[list[str]] + + def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None): credential = {"type": "oauth2", "access_token": access_token} if refresh_token is not None: @@ -5849,7 +5858,12 @@ def test_oauth_authorization_server_404_for_non_bridge_client_forwarded_server() assert exc.value.status_code == 404 -async def _bridge_register_response(server, request_payload, replace_stale_client=False): +async def _bridge_register_response( + server: "MCPServer", + request_payload: _BridgeRegistrationPayload, + replace_stale_client: bool = False, + database_connected: bool = True, +) -> tuple[Response, MagicMock, AsyncMock]: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client_with_server, ) @@ -5879,6 +5893,10 @@ async def _bridge_register_response(server, request_payload, replace_stale_clien new_callable=AsyncMock, return_value=False, ), + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + **({"return_value": MagicMock()} if database_connected else {"side_effect": HTTPException(status_code=500)}), + ), ): response = await register_client_with_server( request=_bridge_mock_request(), @@ -6033,6 +6051,309 @@ async def test_register_bridge_relay_never_persists(): mock_persist.assert_not_called() +_LAPTOP_LOOPBACK_REDIRECT: Final = "http://127.0.0.1:33418/" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type_value", ["oauth_delegate", "true_passthrough"]) +async def test_register_bridge_with_proxy_base_url_registers_gateway_client_not_laptop_loopback( + monkeypatch: pytest.MonkeyPatch, auth_type_value: Literal["oauth_delegate", "true_passthrough"] +) -> None: + """A laptop MCP client (VS Code) registers with its own loopback redirect. With PROXY_BASE_URL set + the gateway must not relay that registration upstream: it registers its one shared client bound + to the gateway /callback, keeps it, and hands the caller a client bound to the caller's loopback. + Relaying the loopback is the customer failure on upstreams that restrict redirect domains.""" + import json + + from litellm.types.mcp import MCPAuth + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + response, mock_async_client, mock_persist = await _bridge_register_response( + _bridge_server(auth_type=MCPAuth(auth_type_value)), + {"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]}, + ) + + posted: Final = mock_async_client.post.call_args.kwargs["json"] + assert posted["redirect_uris"] == ["https://litellm.example.com/callback"], posted + assert posted["client_name"] == "LiteLLM Proxy", posted + mock_persist.assert_awaited_once() + + payload: Final = json.loads(response.body.decode("utf-8")) + assert payload["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT], payload + + +@pytest.mark.asyncio +async def test_register_bridge_without_database_still_relays_client_registration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A database-free gateway cannot keep the shared client past this process, so even with + PROXY_BASE_URL set it relays the caller's own registration instead of stranding the next authorize.""" + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + _, mock_async_client, mock_persist = await _bridge_register_response( + _bridge_server(), + {"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]}, + database_connected=False, + ) + + assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT] + mock_persist.assert_not_called() + + +@pytest.mark.asyncio +async def test_register_bridge_without_proxy_base_url_still_relays_client_registration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Without PROXY_BASE_URL the gateway cannot keep a shared client, so the relay arm remains the + fallback and the caller's own redirect goes upstream as before.""" + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + _, mock_async_client, mock_persist = await _bridge_register_response( + _bridge_server(), + {"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]}, + ) + + assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT] + mock_persist.assert_not_called() + + +def _persisted_shared_gateway_client( + server: "MCPServer", client_id: str = "shared-gateway-client" +) -> LiteLLM_MCPServerTable: + request: Final = _bridge_mock_request() + credentials: Final = { + "client_id": client_id, + "client_secret": None, + "token_endpoint_auth_method": None, + "redirect_uris": [f"{request.base_url.rstrip('/')}/callback"], + "dcr_issuer": server.issuer, + "dcr_server_url": server.url, + } + return LiteLLM_MCPServerTable.model_validate({**server.model_dump(exclude_none=True), "credentials": credentials}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("shared_client_persisted", "expected_redirect_uri", "expected_client_id"), + [ + (True, "https://litellm.example.com/callback", "shared-gateway-client"), + (False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"), + ], +) +async def test_authorize_clientless_bridge_loads_shared_client_before_choosing_the_arm( + monkeypatch: pytest.MonkeyPatch, + shared_client_persisted: bool, + expected_redirect_uri: str, + expected_client_id: str, +) -> None: + """This worker's catalog snapshot has no client_id. When another worker already persisted the shared + client, authorize must load it and send it with the gateway /callback upstream, never the caller's + own client_id or loopback. When no shared client exists yet (a client registered before the gateway + kept one), the caller keeps its own upstream binding through the relay.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server: Final = _bridge_server(url="https://mcp.example.com/mcp") + persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock(return_value=persisted) + mock_update_server: Final = AsyncMock() + with ( + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper", + return_value="mocked_encrypted_state", + ), + ): + response: Final = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="dcr-client-123", + redirect_uri=_LAPTOP_LOOPBACK_REDIRECT, + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + + location: Final = response.headers["location"] + assert location.startswith("https://provider.com/oauth/authorize"), location + query: Final = parse_qs(urlparse(location).query) + assert query["redirect_uri"] == [expected_redirect_uri], query + assert query["client_id"] == [expected_client_id], query + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("shared_client_persisted", "expected_redirect_uri", "expected_client_id"), + [ + (True, "https://litellm.example.com/callback", "shared-gateway-client"), + (False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"), + ], +) +async def test_token_clientless_bridge_loads_shared_client_before_choosing_the_arm( + monkeypatch: pytest.MonkeyPatch, + shared_client_persisted: bool, + expected_redirect_uri: str, + expected_client_id: str, +) -> None: + """The token leg mirrors authorize: with the shared client persisted it redeems the code with that + client and the gateway /callback, and with none yet it relays the caller's own client and redirect.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server: Final = _bridge_server(url="https://mcp.example.com/mcp") + persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock(return_value=persisted) + mock_update_server: Final = AsyncMock() + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma_client), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + ): + data: Final = await _bridge_token_post_data(server, redirect_uri=_LAPTOP_LOOPBACK_REDIRECT) + + assert data["redirect_uri"] == expected_redirect_uri, data + assert data["client_id"] == expected_client_id, data + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dcr_bridge", "expected"), + [(True, "persisted"), (False, "skipped")], +) +async def test_persist_dcr_client_for_client_forwarded_server_only_when_it_is_a_bridge( + dcr_bridge: bool, + expected: Literal["persisted", "skipped"], +) -> None: + """A DCR bridge keeps the shared gateway client on the server row (without stamping oauth2_flow on + a client-forwarded row), while a client-forwarded server without the bridge still stores nothing.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + + server: Final = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=dcr_bridge) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + mock_update: Final = AsyncMock(return_value=MagicMock()) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), + ): + result: Final = await _persist_dcr_client_registration( + server, + {"client_id": "shared-gateway-client", "token_endpoint_auth_method": "none"}, + "https://litellm.example.com/callback", + replace_stale_client=False, + ) + + assert result == expected + if expected == "skipped": + mock_update.assert_not_awaited() + return + written: Final = mock_update.await_args.kwargs["data"] + assert written.oauth2_flow is None, written + assert written.credentials["client_id"] == "shared-gateway-client", written + assert server.client_id == "shared-gateway-client" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("server_has_shared_client", [False, True]) +async def test_token_sealed_ephemeral_client_is_not_replaced_by_persisted_gateway_client( + monkeypatch: pytest.MonkeyPatch, + server_has_shared_client: bool, +) -> None: + """The token leg must use the sealed ephemeral client even if this worker's catalog has a shared client.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server: Final = _bridge_server( + url="https://mcp.example.com/mcp", + client_id="shared-gateway-client" if server_has_shared_client else None, + client_secret="shared-secret" if server_has_shared_client else None, + token_endpoint_auth_method="client_secret_basic" if server_has_shared_client else None, + ) + fake_http_response: Final = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client: Final = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock( + return_value=_persisted_shared_gateway_client(server, client_id="persisted-gateway-client") + ) + mock_update_server: Final = AsyncMock() + with ( + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + ): + await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="upstream-code", + redirect_uri="https://litellm.example.com/callback", + client_id="ephemeral-client", + client_secret="ephemeral-secret" if server_has_shared_client else None, + code_verifier="verifier", + client_token_endpoint_auth_method="client_secret_post" if server_has_shared_client else None, + sealed_client=True, + ) + + data: Final = fake_http_client.post.call_args.kwargs["data"] + assert data["client_id"] == "ephemeral-client", data + assert data["redirect_uri"] == "https://litellm.example.com/callback", data + if server_has_shared_client: + headers: Final = fake_http_client.post.call_args.kwargs["headers"] + assert data["client_secret"] == "ephemeral-secret", data + assert "Authorization" not in headers, headers + + _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index ef95df23ff8..530667fdfdd 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -3384,6 +3384,7 @@ async def test_jwt_bearer_grant_claims_the_jti_before_it_mints(): ("invalid_client", 401), ("invalid_grant", 400), ("invalid_request", 400), + ("invalid_target", 400), ("unsupported_grant_type", 400), ("temporarily_unavailable", 503), ], diff --git a/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py index fdfa0544ad1..f4fa51d57ba 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py @@ -33,6 +33,7 @@ from litellm.proxy._experimental.mcp_server.id_jag_grant import ( jwt_bearer_available, names_mcp_resource, proxy_user_lookup, + resource_names_mcp, select_team, subject_resolver, ) @@ -87,6 +88,16 @@ def _assertion(typ: str = ID_JAG_TYP, **claims: object) -> str: ) +def _unscoped_assertion(**claims: object) -> str: + unscoped_claims: Final = {key: value for key, value in _claims(**claims).items() if key != "scope"} + return jwt.encode( + unscoped_claims, + "unused-by-the-fake-verifier-but-long-enough-for-hs256", + algorithm="HS256", + headers={"typ": ID_JAG_TYP}, + ) + + class _Verifier: """Stands in for the injected ``DecodeAssertion`` port; records each call as a new tuple.""" @@ -714,6 +725,40 @@ def test_the_mcp_resource_audience_is_the_mcp_path_without_the_bare_origin(aud, assert names_mcp_resource(aud, None) is False +@pytest.mark.parametrize( + "claims, expected", + [ + ({}, False), + ({"resource": f"{GATEWAY}/mcp"}, True), + ({"resource": [f"{GATEWAY}/mcp"]}, True), + ({"resource": GATEWAY}, False), + ({"resource": [GATEWAY, f"{GATEWAY}/mcp"]}, False), + ({"resource": "https://other.example"}, None), + ({"resource": [f"{GATEWAY}/mcp", "https://other.example"]}, None), + ({"resource": f"{GATEWAY}/mcp/"}, None), + ({"resource": None}, None), + ({"resource": 7}, None), + ({"resource": []}, None), + ({"resource": [7]}, None), + ], +) +def test_the_signed_resource_claim_must_name_only_gateway_resources( + claims: Mapping[str, object], expected: bool | None +) -> None: + result: Final = resource_names_mcp(claims, GATEWAY) + if isinstance(result, bool): + assert expected is result + return + assert expected is None + assert result.error == "invalid_target" + + +def test_the_signed_resource_claim_is_refused_without_a_gateway_base_url() -> None: + result: Final = resource_names_mcp({"resource": f"{GATEWAY}/mcp"}, None) + assert isinstance(result, SubjectTokenRefusal) + assert result.error == "invalid_target" + + @pytest.mark.asyncio async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_llm_scope(): """End to end through identity_from_id_jag: the same verified claims, with ``aud`` the MCP @@ -736,3 +781,33 @@ async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_l scoped = await _identity(_redemption(_assertion(aud=AUDIENCES[1], scope="litellm:llm")), _Verifier(), _Resolver()) assert isinstance(scoped, AssertionIdentity) assert scoped.scope == "litellm:llm" + + +@pytest.mark.asyncio +async def test_redeeming_for_the_signed_resource_caps_the_identity_before_subject_lookup() -> None: + mcp_claim: Final = _unscoped_assertion(aud=GATEWAY, resource=f"{GATEWAY}/mcp") + narrowed: Final = await _identity(_redemption(mcp_claim), _Verifier(), _Resolver()) + assert isinstance(narrowed, AssertionIdentity) + assert (narrowed.scope, narrowed.allowed_routes) == ("litellm:mcp", ("mcp_inference_routes",)) + + llm_claim: Final = _assertion(aud=GATEWAY, resource=f"{GATEWAY}/mcp", scope="litellm:llm") + refused_scope: Final = await _identity(_redemption(llm_claim), _Verifier(), _Resolver()) + assert isinstance(refused_scope, SubjectTokenRefusal) + assert refused_scope.error == "invalid_scope" + + other_resource: Final = _unscoped_assertion(aud=GATEWAY, resource="https://other.example") + resolver: Final = _Resolver() + refused_resource: Final = await _identity(_redemption(other_resource), _Verifier(), resolver) + assert isinstance(refused_resource, SubjectTokenRefusal) + assert refused_resource.error == "invalid_target" + assert resolver.calls == () + + origin_resource: Final = _unscoped_assertion(aud=GATEWAY, resource=GATEWAY) + full_resource: Final = await _identity(_redemption(origin_resource), _Verifier(), _Resolver()) + assert isinstance(full_resource, AssertionIdentity) + assert (full_resource.scope, full_resource.allowed_routes) == (None, None) + + no_resource: Final = _unscoped_assertion(aud=GATEWAY) + full_control: Final = await _identity(_redemption(no_resource), _Verifier(), _Resolver()) + assert isinstance(full_control, AssertionIdentity) + assert (full_control.scope, full_control.allowed_routes) == (None, None) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 40901d1dd96..d47b2f73a48 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -15634,6 +15634,73 @@ def _mcp_upstream(respond): yield +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("responder", "expects_exchange", "expected_fragments"), + ( + pytest.param( + AsyncMock(return_value=httpx2.Response(500, json={"error": "health-check-boom"})), + True, + ( + "MCP health check failed for", + "upstream exchange:", + "-> HTTP 500", + '"method":"initialize"', + "health-check-boom", + ), + id="upstream-http-500", + ), + pytest.param( + AsyncMock(side_effect=httpx2.ConnectError("connection failed")), + False, + ("MCP health check failed for",), + id="connect-error", + ), + ), +) +async def test_health_check_logs_upstream_exchange_only_for_http_failures( + caplog: pytest.LogCaptureFixture, + responder: AsyncMock, + expects_exchange: bool, + expected_fragments: tuple[str, ...], +) -> None: + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(responder)) + token: Final = "health-check-secret-0123456789" + server: Final = MCPServer( + server_id="health-check-upstream", + name="health_check_upstream", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=token, + ) + manager.registry[server.server_id] = server + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = await manager.health_check_server(server.server_id) + + assert result.status == "unhealthy" + warning_messages: Final = tuple( + record.getMessage() + for record in caplog.records + if record.name == "LiteLLM" + and record.levelno == logging.WARNING + and "MCP health check failed for" in record.getMessage() + ) + assert len(warning_messages) == 1 + warning: Final = warning_messages[0] + assert all(fragment in warning for fragment in expected_fragments) + assert ("upstream exchange:" in warning) is expects_exchange + assert not any( + record.name == "LiteLLM" + and record.levelno == logging.WARNING + and "run_with_session failed" in record.getMessage() + for record in caplog.records + ) + assert token not in caplog.text + assert responder.await_args_list[0].args[0].headers["Authorization"] == f"Bearer {token}" + + class _DiscoveryUpstream: def __init__(self) -> None: self.requests: tuple[tuple[str, str], ...] = () @@ -19989,6 +20056,74 @@ async def test_disabled_stdio_tools_are_empty_initially_and_reject_continuation( assert page.next_cursor is None +@pytest.mark.asyncio +async def test_aggregate_bounds_each_paginated_server_by_its_own_listing_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import JSONRPCRequest, PaginatedRequestParams + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import catalog, operations + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix + + monkeypatch.setenv("LITELLM_SALT_KEY", "paginated-listing-deadline-test") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(manager_module, "MCP_TOOL_LISTING_TIMEOUT", 0.0) + extended: Final = MCPServer( + server_id="extended", + name="extended", + transport=MCPTransport.http, + url="https://extended.example/mcp", + timeout=60.0, + ) + global_deadline: Final = MCPServer( + server_id="global", name="global", transport=MCPTransport.http, url="https://global.example/mcp" + ) + servers: Final = (extended, global_deadline) + manager: Final = MCPServerManager() + manager.registry = {server.server_id: server for server in servers} + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + context: Final = operations.prepare_context(UserAPIKeyAuth(api_key="owned-caller", user_id="alice")) + + async def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = _JSONRPC_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + result: Final = { + "initialize": { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "test", "version": "1"}, + "capabilities": {"tools": {}}, + }, + "server/discover": { + "supportedVersions": ["2026-07-28"], + "capabilities": {"tools": {}}, + "ttlMs": 60000, + "cacheScope": "private", + "resultType": "complete", + }, + "tools/list": { + "tools": [{"name": "echo", "inputSchema": {"type": "object"}}], + "ttlMs": 60000, + "cacheScope": "private", + "resultType": "complete", + }, + }[payload.method] + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + with _mcp_upstream(respond): + listing: Final = await catalog.aggregate_gateway_tools( + context, PaginatedRequestParams(), servers, {}, record_listing=True, enforce_rate_limits=False + ) + + assert [tool.name for tool in listing.tools] == [f"{get_server_prefix(extended)}{MCP_TOOL_PREFIX_SEPARATOR}echo"] + assert listing.outcomes["extended"].tag == "ok" + assert listing.outcomes["global"].tag == "timeout" + + @pytest.mark.asyncio async def test_failed_aggregate_continuation_preserves_only_delivered_tool_metadata(monkeypatch): from mcp.shared.exceptions import MCPError diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 21958e88008..7bdcb8afd37 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -4782,6 +4782,51 @@ def test_view_only_user_denied_credential_connection_routes(method: str, route: _internal_user_check(route, method, role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) +_SELF_SCOPED_CREDENTIAL_WRITES = ( + ("POST", "/credentials/x/user_connection/start"), + ("POST", "/credentials/x/user_connection/poll"), + ("DELETE", "/credentials/x/user_connection"), + ("POST", "/v1/mcp/server/oauth/srv-1/register"), + ("POST", "/v1/mcp/server/srv-1/oauth-user-credential"), + ("DELETE", "/v1/mcp/server/srv-1/oauth-user-credential"), + ("POST", "/v1/mcp/server/srv-1/user-credential"), + ("DELETE", "/v1/mcp/server/srv-1/user-credential"), + ("POST", "/v1/mcp/server/srv-1/user-env-vars"), + ("DELETE", "/v1/mcp/server/srv-1/user-env-vars"), +) + + +@pytest.mark.parametrize(("method", "route"), _SELF_SCOPED_CREDENTIAL_WRITES) +def test_proxy_admin_viewer_allowed_self_scoped_credential_writes(method: str, route: str) -> None: + """Regression for LIT-9320: these handlers only write the caller's own credential + row, so a read-only admin can use them like /user/password/change.""" + assert _internal_user_check(route, method, role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) is None + + +@pytest.mark.parametrize( + ("method", "route"), + [ + ("POST", "/v1/mcp/server"), + ("PUT", "/v1/mcp/server"), + ("DELETE", "/v1/mcp/server/srv-1"), + ("POST", "/v1/mcp/server/import"), + ("POST", "/v1/mcp/server/srv-1/pin"), + ("PATCH", "/v1/mcp/server/srv-1/oauth-user-credential"), + ("PUT", "/v1/mcp/server/srv-1/user-credential"), + ("PATCH", "/v1/mcp/server/srv-1/user-env-vars"), + ("POST", "/v1/mcp/server/srv-1/user-credential/rotate"), + ("POST", "/v1/mcp/server/oauth/srv-1/token"), + ("DELETE", "/v1/mcp/server/oauth/srv-1/register"), + ("PATCH", "/credentials/x/user_connection"), + ("DELETE", "/credentials/user_connections"), + ("POST", "/key/generate"), + ], +) +def test_proxy_admin_viewer_denied_mcp_admin_and_sibling_writes(method: str, route: str) -> None: + with pytest.raises(HTTPException, match="user not allowed to access this route"): + _internal_user_check(route, method, role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + + def test_user_connection_route_registration_order_wins_over_generic_credential_routes() -> None: """The :path placeholders in the generic CRUD routes could swallow the user_connection paths; the concrete routes must register (and match) first.""" diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index f6f1c2fec3b..7be98521f8e 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -28,6 +28,7 @@ from litellm.proxy.client.cli.commands.agents import ( run_agent, verify_proxy_key, ) +from litellm.proxy.client.cli.main import cli as lite AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" @@ -1490,6 +1491,13 @@ class TestAgentCommands: def test_one_command_per_known_agent(self): assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode", "pi"} + @pytest.mark.parametrize("agent", [c.name for c in agent_commands()]) + def test_lite_runs_every_known_agent_command(self, agent): + result = self.runner.invoke(lite, [agent, "--help"], prog_name="lite") + + assert result.exit_code == 0, result.output + assert result.output.splitlines()[0].startswith(f"Usage: lite {agent} ") + def test_pi_is_hidden_from_help_but_still_registered(self): hidden_by_name = {c.name: c.hidden for c in agent_commands()} assert hidden_by_name == {"claude": False, "codex": False, "opencode": False, "pi": True} diff --git a/tests/unit/proxy/client/cli/test_claude_settings.py b/tests/unit/proxy/client/cli/test_claude_settings.py index 9353c149d15..9451598aa40 100644 --- a/tests/unit/proxy/client/cli/test_claude_settings.py +++ b/tests/unit/proxy/client/cli/test_claude_settings.py @@ -15,6 +15,7 @@ import pytest from click.testing import CliRunner from litellm.litellm_core_utils.private_json import commit_staged_json, write_private_bytes +from litellm.proxy.client.cli.commands.autoroute.process import AUTOROUTE_DIR from litellm.proxy.client.cli.commands.claude_settings import ( ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, AUTOROUTE_BACKUP_PATH, @@ -196,8 +197,6 @@ class TestConflictingOwnersOfTheSettingsFile: def test_the_registry_matches_the_paths_the_commands_actually_use(self): """A second definition of the autoroute dir must not drift from this one.""" - from litellm.proxy.client.cli.commands.autoroute.process import AUTOROUTE_DIR - assert AUTOROUTE_BACKUP_PATH == AUTOROUTE_DIR / "claude_settings_backup.json" assert {o.backup_path for o in SETTINGS_FILE_OWNERS} == {BACKUP_PATH, AUTOROUTE_BACKUP_PATH} assert {o.stop_command for o in SETTINGS_FILE_OWNERS} == {"lite down", "lite autoroute stop"} diff --git a/tests/unit/proxy/client/cli/test_credentials_commands.py b/tests/unit/proxy/client/cli/test_credentials_commands.py index fe9bc333f37..cdbd73057a0 100644 --- a/tests/unit/proxy/client/cli/test_credentials_commands.py +++ b/tests/unit/proxy/client/cli/test_credentials_commands.py @@ -14,10 +14,6 @@ from litellm.proxy.client.cli.main import cli def mock_credentials_client(monkeypatch): """Patch the CredentialsManagementClient used by the CLI commands.""" mock_client = MagicMock() - monkeypatch.setattr( - "litellm.proxy.client.credentials.CredentialsManagementClient", - mock_client, - ) monkeypatch.setattr( "litellm.proxy.client.cli.commands.credentials.CredentialsManagementClient", mock_client, diff --git a/tests/unit/proxy/client/cli/test_global_options.py b/tests/unit/proxy/client/cli/test_global_options.py index 3b9f6820b69..8a1e9a3d141 100644 --- a/tests/unit/proxy/client/cli/test_global_options.py +++ b/tests/unit/proxy/client/cli/test_global_options.py @@ -473,3 +473,13 @@ def test_version_command_without_a_context_object_fails_reading_it(cli_runner): assert result.exit_code == 1 assert type(result.exception) is AttributeError assert result.output == "" + + +def test_every_listed_command_runs_the_command_of_that_name(): + from litellm.proxy.client.cli.commands.agents import agent_commands + + commands: Final = {name: cli.commands[name] for name in cli.commands} + + assert {name: command.name for name, command in commands.items()} == {name: name for name in commands} + assert {command.name for command in agent_commands()} <= commands.keys() + assert all(cli.commands[name] is command for name, command in commands.items()) diff --git a/tests/unit/proxy/common_utils/test_slow_gc_logger.py b/tests/unit/proxy/common_utils/test_slow_gc_logger.py new file mode 100644 index 00000000000..4bc78789b59 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_slow_gc_logger.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import gc +import re +from itertools import count +from typing import Final + +from litellm.proxy.common_utils.slow_gc_logger import SlowGcLogger + + +def _collect_once(step_seconds: float) -> tuple[str, ...]: + ticks: Final = count(step=step_seconds) + lines: Final[list[str]] = [] # mutable-ok: collects what the logger emits during one real gc.collect() + logger: Final = SlowGcLogger(threshold_seconds=0.1, log=lines.append, clock=lambda: next(ticks)) + gc.callbacks.append(logger) + try: + gc.collect() + finally: + gc.callbacks.remove(logger) + return tuple(lines) + + +def test_a_full_collection_past_the_threshold_is_logged_with_its_generation_and_length() -> None: + lines: Final = _collect_once(step_seconds=0.25) + assert any(re.fullmatch(r"Garbage collection gen=2 collected=\d+ took 0\.250 seconds", line) for line in lines), ( + lines + ) + + +def test_a_collection_under_the_threshold_is_not_logged() -> None: + assert _collect_once(step_seconds=0.01) == () diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 99b500ec315..4b776bfbb1a 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -1,4 +1,5 @@ import asyncio +import importlib.resources import json import os import signal @@ -202,6 +203,19 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client( DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"] +ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV = [ + "db", + "execute", + "--schema", + "schema.prisma", + "--file", + str( + importlib.resources.files("litellm_proxy_extras") + / "migrations" + / "20261015120000_add_environment_variables_revision_trigger" + / "migration.sql" + ), +] def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url): @@ -221,7 +235,7 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_ assert PrismaManager.setup_database(use_migrate=False) is True - assert fake_prisma_cli.calls == [DB_PUSH_ARGV] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV] assert applied == [True] @@ -254,7 +268,7 @@ def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_p ) assert PrismaManager.setup_database(use_migrate=False) is True - assert fake_prisma_cli.calls == [DB_PUSH_ARGV] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV] def _entra_jwt(expires_in_seconds: int) -> str: @@ -447,7 +461,7 @@ def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_d monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") assert PrismaManager.setup_database(use_migrate=False) is True - assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV, ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV] assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py index 180cdbe5bb5..b145b990539 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py @@ -1,3 +1,4 @@ +import asyncio import os from unittest.mock import AsyncMock, MagicMock, patch @@ -534,6 +535,33 @@ class TestNomaV2ActionBehavior: assert result == inputs + @pytest.mark.asyncio + async def test_a_scan_cancelled_before_noma_answers_is_recorded_as_not_run(self, noma_v2_guardrail): + scan_sent = asyncio.Event() + + async def noma_never_answers(payload): + scan_sent.set() + await asyncio.Event().wait() + + request_data = {"metadata": {}} + with patch.object(noma_v2_guardrail, "_call_noma_scan", noma_never_answers): + scan = asyncio.ensure_future( + noma_v2_guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data=request_data, + input_type="response", + ) + ) + await scan_sent.wait() + scan.cancel() + with pytest.raises(asyncio.CancelledError): + await scan + + assert [ + record["guardrail_status"] + for record in request_data["metadata"]["standard_logging_guardrail_information"] + ] == ["not_run"] + @pytest.mark.asyncio async def test_fail_closed_on_technical_scan_failure_when_block_failures_true(self): guardrail = NomaV2Guardrail( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 7d0c81af4c2..bd1e9ae4068 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3789,6 +3789,246 @@ async def test_output_parse_pii_numbered_tokens_across_chunks(): assert pii_tokens[""] == CHUNK_MARKER_TWO +@pytest.mark.parametrize( + ("text", "analyze_results", "expected_text", "expected_tokens", "expected_counts"), + [ + pytest.param( + "Repeat exactly: my number is 4251234567", + [ + {"entity_type": "PHONE_NUMBER", "start": 29, "end": 39, "score": 0.75}, + {"entity_type": "US_BANK_NUMBER", "start": 29, "end": 39, "score": 0.05}, + {"entity_type": "US_DRIVER_LICENSE", "start": 29, "end": 39, "score": 0.01}, + ], + "Repeat exactly: my number is ", + {"": "4251234567"}, + {"PHONE_NUMBER": 1}, + id="same-span", + ), + pytest.param( + "contact john.smith@example.com today", + [ + {"entity_type": "EMAIL_ADDRESS", "start": 8, "end": 30, "score": 1.0}, + {"entity_type": "DATE_TIME", "start": 31, "end": 36, "score": 0.85}, + {"entity_type": "URL", "start": 8, "end": 15, "score": 0.5}, + {"entity_type": "URL", "start": 19, "end": 30, "score": 0.5}, + ], + "contact ", + {"": "john.smith@example.com", "": "today"}, + {"EMAIL_ADDRESS": 1, "DATE_TIME": 1}, + id="nested", + ), + pytest.param( + "her phone is +33 1 42 68 53 00.", + [ + {"entity_type": "DATE_TIME", "start": 25, "end": 30, "score": 0.85}, + {"entity_type": "PHONE_NUMBER", "start": 13, "end": 30, "score": 0.75}, + ], + "her phone is .", + {"": "+33 1 42 68 53 00"}, + {"PHONE_NUMBER": 1}, + id="container-outranks-higher-scored-fragment", + ), + pytest.param( + "ids ab0123456789 and cd0123456789", + [ + {"entity_type": "AAA", "start": 4, "end": 16, "score": 0.5}, + {"entity_type": "BBB", "start": 21, "end": 29, "score": 0.5}, + {"entity_type": "CCC", "start": 25, "end": 33, "score": 0.9}, + ], + "ids and ", + {"": "ab0123456789", "": "cd0123456789"}, + {"AAA": 1, "CCC": 1}, + id="partial-overlap-in-second-run", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 15, "score": 0.5}, + {"entity_type": "BBB", "start": 3, "end": 8, "score": 0.9}, + ], + "id ", + {"": "ab0123456789"}, + {"AAA": 1}, + id="container-sharing-start-outranks-fragment", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 10, "score": 0.9}, + {"entity_type": "BBB", "start": 7, "end": 15, "score": 0.5}, + ], + "id ", + {"": "ab0123456789"}, + {"AAA": 1}, + id="earlier-higher-scored-detection-labels-the-whole-run", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 6, "end": 15, "score": 0.5}, + {"entity_type": "BBB", "start": 3, "end": 8, "score": 0.5}, + ], + "id ", + {"": "ab0123456789"}, + {"AAA": 1}, + id="score-tie-goes-to-the-longer-detection", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 9, "score": 0.5}, + {"entity_type": "BBB", "start": 9, "end": 15, "score": 0.5}, + ], + "id ", + {"": "ab0123", "": "456789"}, + {"AAA": 1, "BBB": 1}, + id="touching-detections-stay-separate", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 15, "score": 0.5}, + {"entity_type": "BBB", "start": 3, "end": 15, "score": 0.5}, + ], + "id ", + {"": "ab0123456789"}, + {"BBB": 1}, + id="identical-spans-with-equal-scores-label-the-same-in-any-order", + ), + pytest.param( + "My name is John Smith, phone +1 415 555 2671, please check the order.", + [ + {"entity_type": "EMPTY", "start": 25, "end": 25, "score": 0.99}, + {"entity_type": "INVERTED", "start": 40, "end": 30, "score": 0.99}, + {"entity_type": "PERSON", "start": 11, "end": 21, "score": 0.9}, + {"entity_type": "PHONE_NUMBER", "start": 29, "end": 44, "score": 0.85}, + {"entity_type": "UK_NHS", "start": 32, "end": 44, "score": 0.6}, + ], + "My name is , phone , please check the order.", + {"": "John Smith", "": "+1 415 555 2671"}, + {"PERSON": 1, "PHONE_NUMBER": 1}, + id="empty-and-inverted-spans-are-ignored", + ), + pytest.param( + "Code ABC-12345-XYZ is valid.", + [ + {"entity_type": "GENERIC_ID", "start": 5, "end": 14, "score": 0.5}, + {"entity_type": "SECRET_KEY", "start": 9, "end": 18, "score": 0.95}, + ], + "Code is valid.", + {"": "ABC-12345-XYZ"}, + {"SECRET_KEY": 1}, + id="later-higher-scored-detection-labels-the-whole-run", + ), + ], +) +@pytest.mark.parametrize("order", ["analyzer", "reversed"]) +def test_numbered_tokens_collapse_overlapping_detections_into_one_token_per_run( + text: str, + analyze_results: list[PresidioAnalyzeResponseItem], + expected_text: str, + expected_tokens: dict[str, str], + expected_counts: dict[str, int], + order: Literal["analyzer", "reversed"], +) -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking(mock_testing=True) + request_data: Final = {"metadata": {}} + masked_entity_count: Final = {} + ordered: Final = analyze_results if order == "analyzer" else analyze_results[::-1] + + masked: Final = guardrail._finalize_presidio_anonymize_numbered_tokens( + text, ordered, request_data, masked_entity_count + ) + + assert masked == expected_text + assert request_data["metadata"]["pii_tokens"] == expected_tokens + assert masked_entity_count == expected_counts + assert OPTIONAL_PresidioPIIMasking._unmask_pii_text(masked, expected_tokens) == text + + +@pytest.mark.parametrize( + ("parts", "expected_masked", "expected_tokens"), + [ + pytest.param( + ( + ( + "My number is 4251234567.", + [ + {"entity_type": "PHONE_NUMBER", "start": 13, "end": 23, "score": 0.75}, + {"entity_type": "US_BANK_NUMBER", "start": 13, "end": 23, "score": 0.05}, + {"entity_type": "US_DRIVER_LICENSE", "start": 13, "end": 23, "score": 0.01}, + ], + ), + ( + "Also call +1 (206) 555-0199.", + [{"entity_type": "PHONE_NUMBER", "start": 10, "end": 27, "score": 0.4}], + ), + ), + ("My number is .", "Also call ."), + {"": "4251234567", "": "+1 (206) 555-0199"}, + id="overlapping-run-then-plain-detection", + ), + pytest.param( + ( + ("Call John Smith.", [{"entity_type": "PERSON", "start": 5, "end": 15, "score": 0.85}]), + ("Again John Smith.", [{"entity_type": "PERSON", "start": 6, "end": 16, "score": 0.85}]), + ), + ("Call .", "Again ."), + {"": "John Smith"}, + id="same-value-keeps-its-token", + ), + pytest.param( + ( + ("Call John Smith.", [{"entity_type": "PERSON", "start": 5, "end": 15, "score": 0.85}]), + ( + "Ann Lee and Bob Ray.", + [ + {"entity_type": "PERSON", "start": 0, "end": 7, "score": 0.85}, + {"entity_type": "PERSON", "start": 12, "end": 19, "score": 0.85}, + ], + ), + ), + ("Call .", " and ."), + {"": "John Smith", "": "Bob Ray", "": "Ann Lee"}, + id="bump-skips-every-claimed-number", + ), + pytest.param( + ( + ( + "Ann Lee and Bob Ray.", + [ + {"entity_type": "PERSON", "start": 0, "end": 7, "score": 0.85}, + {"entity_type": "PERSON", "start": 12, "end": 19, "score": 0.85}, + ], + ), + ("Call John Smith.", [{"entity_type": "PERSON", "start": 5, "end": 15, "score": 0.85}]), + ), + (" and .", "Call ."), + {"": "Ann Lee", "": "Bob Ray", "": "John Smith"}, + id="spare-number-starts-above-every-earlier-part", + ), + ], +) +def test_numbered_tokens_stay_unique_across_parts_of_one_request( + parts: tuple[tuple[str, list[PresidioAnalyzeResponseItem]], ...], + expected_masked: tuple[str, ...], + expected_tokens: dict[str, str], +) -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking(mock_testing=True) + request_data: Final = {"metadata": {}} + + masked: Final = tuple( + guardrail._finalize_presidio_anonymize_numbered_tokens(text, results, request_data, {}) + for text, results in parts + ) + + assert masked == expected_masked + assert request_data["metadata"]["pii_tokens"] == expected_tokens + assert OPTIONAL_PresidioPIIMasking._unmask_pii_text(" | ".join(masked), expected_tokens) == " | ".join( + text for text, _ in parts + ) + + @pytest.mark.asyncio async def test_analyze_text_chunked_failure_stays_fail_closed(): """If one chunk still fails, the chunked path raises exactly like a single diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 4e19c7ed328..2ee4c8c72c0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) import json import logging import re +from typing import Final from unittest.mock import patch import pytest @@ -1379,6 +1380,54 @@ class TestToolPermissionGuardrailAnthropicMessages: assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served" + @pytest.mark.asyncio + async def test_rewrite_mode_keeps_the_allowed_tool_use_verdict_with_its_explanation(self): + """Claude Code auto mode reads the classifier verdict from message_delta. The rewrite changes + no text the explanation could repeat, so the allowed call keeps its verdict whole.""" + default_on_rewriting: Final = ToolPermissionGuardrail( + guardrail_name="anthropic-rewrite-default-on", + rules=self.rules, + default_action="deny", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + bash_verdict: Final = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"} + safeguard_results: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"tu_ok": bash_verdict, "tu_bad": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + bash_chunks: Final = self._sse_chunks("Bash", "tu_ok") + read_chunks: Final = self._sse_chunks("Read", "tu_bad") + read_tool_use: Final = [chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7]] + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 5}, + } + chunks: Final = [ + *bash_chunks[:7], + *read_tool_use, + f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(), + bash_chunks[-1], + ] + + body: Final = b"".join(await self._drain(default_on_rewriting, chunks)).decode() + delta: Final = next( + json.loads(line[6:])["delta"] + for line in body.splitlines() + if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_delta" + ) + assert "Permission denied" in body + assert delta["safeguard_results"] == [ + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {"tu_ok": bash_verdict}}} + ] + @pytest.mark.asyncio async def test_message_start_without_a_dict_message_fails_closed(self): """Malformed SSE must not be forwarded unscanned. diff --git a/tests/unit/proxy/guardrails/test_anthropic_sse.py b/tests/unit/proxy/guardrails/test_anthropic_sse.py new file mode 100644 index 00000000000..95f596aab11 --- /dev/null +++ b/tests/unit/proxy/guardrails/test_anthropic_sse.py @@ -0,0 +1,316 @@ +import json +from collections.abc import Callable +from itertools import chain +from typing import Final + +import pytest + +from litellm.proxy.guardrails.anthropic_sse import ( + anthropic_sse_chunks_from_response, + assemble_anthropic_sse_stream, +) +from litellm.types.utils import Choices, ModelResponse + + +def _frame(event: str, data: dict[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +def _decoded_event(chunk: bytes) -> tuple[str, dict[str, object]]: + event_line, data_line = chunk.decode().split("\n")[:2] + assert event_line.startswith("event: ") + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +def _decode(chunks: tuple[bytes, ...]) -> tuple[tuple[str, dict[str, object]], ...]: + return tuple(_decoded_event(chunk) for chunk in chunks) + + +def _rewritten_delta( + upstream: tuple[bytes, ...], rewrite: Callable[[ModelResponse], None], keep_explanations: bool = False +) -> object: + assembled: Final = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + rewrite(assembled) + rewritten: Final = anthropic_sse_chunks_from_response(assembled, upstream, keep_explanations=keep_explanations) + return next(payload["delta"] for event, payload in _decode(rewritten) if event == "message_delta") + + +def _set_text(text: str) -> Callable[[ModelResponse], None]: + def rewrite(assembled: ModelResponse) -> None: + choice: Final = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = text + + return rewrite + + +def _tool_use_frames(index: int, tool_use_id: str, name: str) -> tuple[bytes, ...]: + return ( + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": index, + "content_block": {"type": "tool_use", "id": tool_use_id, "name": name, "input": {}}, + }, + ), + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "input_json_delta", "partial_json": '{"command": "ls"}'}, + }, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": index}), + ) + + +def _upstream_stream( + delta: dict[str, object], text: str = "call 555-0100", tool_uses: tuple[tuple[str, str], ...] = () +) -> tuple[bytes, ...]: + tool_use_frames = tuple( + chain.from_iterable( + _tool_use_frames(index, tool_use_id, name) for index, (tool_use_id, name) in enumerate(tool_uses, start=1) + ) + ) + return ( + _frame( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_upstream_1", + "type": "message", + "role": "assistant", + "model": "claude-upstream", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + }, + ), + _frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + *tool_use_frames, + _frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 2}}), + _frame("message_stop", {"type": "message_stop"}), + ) + + +def _safeguard_results(tool_uses: dict[str, dict[str, str]]) -> list[dict[str, object]]: + return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_uses}}] + + +SAFEGUARD_RESULTS: Final = _safeguard_results({}) +REFUSAL_STOP_DETAILS: Final = {"type": "refusal", "category": "cyber", "explanation": None} + + +@pytest.mark.parametrize( + ("stop_reason", "delta_fields"), + [ + ("end_turn", {"safeguard_results": SAFEGUARD_RESULTS}), + ("refusal", {"stop_details": REFUSAL_STOP_DETAILS}), + ], + ids=["safeguard_results", "refusal_stop_details"], +) +def test_rewritten_stream_keeps_the_upstream_message_delta_fields(stop_reason: str, delta_fields: dict[str, object]): + """A guardrail that rewrites a streamed response assembles it into the OpenAI shape, which has no + slot for ``safeguard_results`` or ``stop_details``, so the re-emitted ``message_delta`` used to + drop the classifier verdict a Claude Code auto-mode session reads from it.""" + upstream_delta: dict[str, object] = {"stop_reason": stop_reason, "stop_sequence": None, **delta_fields} + upstream = _upstream_stream(upstream_delta) + assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = "call " + + decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream)) + payloads = {event: payload for event, payload in decoded} + + assert payloads["content_block_delta"]["delta"] == {"type": "text_delta", "text": "call "} + assert payloads["message_delta"]["delta"] == upstream_delta + assert not set(delta_fields) & set(payloads["message_start"]["message"]) + + +SSN: Final = "123-45-6789" + + +@pytest.mark.parametrize( + ("upstream_delta", "tool_uses", "expected_fields"), + [ + ( + { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": f"Request involves SSN {SSN}"}, + }, + (), + {"stop_details": {"type": "refusal", "category": "cyber", "explanation": None}}, + ), + ( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": f"leaks SSN {SSN}"}} + ), + }, + (("toolu_bash", "Bash"),), + {"safeguard_results": _safeguard_results({"toolu_bash": {"type": "evaluated", "outcome": "flagged"}})}, + ), + ( + { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": [ + { + "type": "prompt_injection", + "status": {"type": "available", "explanation": f"quotes SSN {SSN}"}, + "explanation": f"saw SSN {SSN}", + } + ], + }, + (), + {"safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}]}, + ), + ], + ids=["stop_details_explanation", "verdict_explanation", "unknown_shape_explanation"], +) +def test_masked_stream_drops_the_unmasked_upstream_explanations( + upstream_delta: dict[str, object], tool_uses: tuple[tuple[str, str], ...], expected_fields: dict[str, object] +): + upstream = _upstream_stream(upstream_delta, text=f"my SSN is {SSN}", tool_uses=tool_uses) + assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = "my SSN is " + + rewritten = anthropic_sse_chunks_from_response(assembled, upstream) + message_delta = next(payload for event, payload in _decode(rewritten) if event == "message_delta") + + assert not [chunk for chunk in rewritten if SSN.encode() in chunk] + assert {key: message_delta["delta"][key] for key in expected_fields} == expected_fields + + +def test_masked_stream_drops_safeguard_results_it_cannot_read(): + upstream_delta: Final = { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": {"type": "dangerous_tool_use", "explanation": f"saw SSN {SSN}"}, + } + + assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("my SSN is ")) == { + "stop_reason": "end_turn", + "stop_sequence": None, + } + + +def test_rewritten_stream_keeps_only_the_verdicts_of_tool_uses_it_still_delivers(): + kept_verdict = {"type": "evaluated", "outcome": "not_flagged"} + upstream = _upstream_stream( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_glob": {"type": "evaluated", "outcome": "flagged"}, "toolu_bash": kept_verdict} + ), + }, + tool_uses=(("toolu_glob", "Glob"), ("toolu_bash", "Bash")), + ) + assembled = assemble_anthropic_sse_stream(upstream) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.tool_calls = [call for call in choice.message.tool_calls or [] if call.id == "toolu_bash"] + choice.message.content = "Permission denied: Tool 'Glob' denied" + + decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream)) + message_delta = next(payload for event, payload in decoded if event == "message_delta") + delivered_tool_use_ids = [ + payload["content_block"]["id"] + for event, payload in decoded + if event == "content_block_start" and payload["content_block"]["type"] == "tool_use" + ] + + assert delivered_tool_use_ids == ["toolu_bash"] + assert message_delta["delta"]["safeguard_results"] == _safeguard_results({"toolu_bash": kept_verdict}) + + +def test_rewritten_refusal_gets_no_stop_details_the_upstream_did_not_send(): + """The OpenAI round trip sets ``stop_details`` on every refusal, and the repo reads a refusal + carrying one as a safeguard refusal, so a plain refusal must not come back as one.""" + upstream_delta: Final = {"stop_reason": "refusal", "stop_sequence": None} + + assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("call ")) == upstream_delta + + +def test_rewritten_stream_keeps_safeguard_shapes_it_does_not_know(): + upstream_delta: Final = { + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": {"type": "pause_turn", "category": None, "explanation": None, "policy_url": "https://x.test"}, + "safeguard_results": [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"toolu_bash": {"type": "unavailable", "reason": "overloaded"}}, + }, + "classifier_version": "2026-10", + }, + {"type": "prompt_injection", "status": {"type": "skipped"}}, + ], + } + upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),)) + + assert _rewritten_delta(upstream, _set_text("call ")) == upstream_delta + + +def test_rewrite_that_changes_no_text_keeps_the_upstream_explanations(): + upstream_delta: Final = { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}} + ), + } + upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),)) + + assert _rewritten_delta(upstream, _set_text("Permission denied"), keep_explanations=True) == upstream_delta + + +def test_rewritten_stream_drops_the_verdict_of_a_tool_use_whose_input_the_guardrail_changed(): + """A verdict scored on the upstream input must not vouch for the input the client runs.""" + verdict: Final = {"type": "evaluated", "outcome": "not_flagged"} + upstream: Final = _upstream_stream( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results({"toolu_masked": verdict, "toolu_untouched": verdict}), + }, + tool_uses=(("toolu_masked", "Bash"), ("toolu_untouched", "Bash")), + ) + + def mask_first_tool_input(assembled: ModelResponse) -> None: + choice: Final = assembled.choices[0] + assert isinstance(choice, Choices) + tool_calls: Final = choice.message.tool_calls or [] + assert [call.id for call in tool_calls] == ["toolu_masked", "toolu_untouched"] + tool_calls[0].function.arguments = '{"command": ""}' + + rewritten_delta: Final = _rewritten_delta(upstream, mask_first_tool_input) + + assert isinstance(rewritten_delta, dict) + assert rewritten_delta["safeguard_results"] == _safeguard_results({"toolu_untouched": verdict}) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 57e0a4fc47c..a5ebf41eebd 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -9584,6 +9584,97 @@ async def test_rotate_master_key_reencrypts_model_params_in_place( ), "api_key must be stored re-encrypted under the new master key, not in plaintext" +@pytest.mark.asyncio +async def test_rotate_master_key_keeps_environment_variables_it_cannot_decrypt(monkeypatch): + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper + from litellm.proxy.management_endpoints.key_management_endpoints import rotate_master_key + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") + under_another_key = encrypt_value_helper("stored-under-another-key", new_encryption_key="sk-another-master-key") + for key in ("ROTATION_TEST_ENCRYPTED", "ROTATION_TEST_PLAINTEXT", "ROTATION_TEST_UNDER_ANOTHER_KEY"): + monkeypatch.setenv(key, "restored-by-monkeypatch") + monkeypatch.delenv(key) + prisma_client = AsyncMock() + prisma_client.db = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_config.update = AsyncMock() + prisma_client.writer_db.litellm_config.find_many = AsyncMock( + return_value=[ + SimpleNamespace( + param_name="environment_variables", + param_value={ + "ROTATION_TEST_ENCRYPTED": encrypt_value_helper("stored-encrypted"), + "ROTATION_TEST_PLAINTEXT": "stored-plaintext", + "ROTATION_TEST_UNDER_ANOTHER_KEY": under_another_key, + }, + ) + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", ProxyConfig()) + + await rotate_master_key( + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-9876"), + current_master_key="sk-old-master-key", + new_master_key="sk-new-master-key", + ) + + stored = prisma_client.db.litellm_config.update.call_args.kwargs["data"]["param_value"].data + rotated = decrypt_if_encrypted_with(stored["ROTATION_TEST_ENCRYPTED"], "sk-new-master-key") + assert {**stored, "ROTATION_TEST_ENCRYPTED": rotated} == { + "ROTATION_TEST_ENCRYPTED": "stored-encrypted", + "ROTATION_TEST_PLAINTEXT": "stored-plaintext", + "ROTATION_TEST_UNDER_ANOTHER_KEY": under_another_key, + } + + +@pytest.mark.asyncio +async def test_rotate_master_key_keeps_an_environment_variable_the_read_replica_has_not_caught_up_on(monkeypatch): + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper + from litellm.proxy.management_endpoints.key_management_endpoints import rotate_master_key + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") + monkeypatch.setenv("ROTATION_TEST_SAVED_A_MOMENT_AGO", "restored-by-monkeypatch") + monkeypatch.delenv("ROTATION_TEST_SAVED_A_MOMENT_AGO") + + def rows(values: dict[str, str]) -> list[SimpleNamespace]: + return [SimpleNamespace(param_name="environment_variables", param_value=values)] + + prisma_client = AsyncMock() + prisma_client.db = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_config.update = AsyncMock() + prisma_client.db.litellm_config.find_many = AsyncMock(return_value=rows({})) + prisma_client.writer_db.litellm_config.find_many = AsyncMock( + return_value=rows({"ROTATION_TEST_SAVED_A_MOMENT_AGO": encrypt_value_helper("saved-a-moment-ago")}) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", ProxyConfig()) + + await rotate_master_key( + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-9876"), + current_master_key="sk-old-master-key", + new_master_key="sk-new-master-key", + ) + + stored = prisma_client.db.litellm_config.update.call_args.kwargs["data"]["param_value"].data + assert {key: decrypt_if_encrypted_with(value, "sk-new-master-key") for key, value in stored.items()} == { + "ROTATION_TEST_SAVED_A_MOMENT_AGO": "saved-a-moment-ago" + } + + async def test_default_key_generate_params_duration(monkeypatch): """ Test that default_key_generate_params with 'duration' is applied diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d59c655835..65c4a8faeb7 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -3545,6 +3545,95 @@ class TestTemporaryMCPSessionEndpoints: assert exchange_mock.await_args.kwargs["redirect_uri"] == "https://litellm.example.com/callback" assert exchange_mock.await_args.kwargs["client_token_endpoint_auth_method"] == "client_secret_basic" + @pytest.mark.asyncio + async def test_mcp_token_prefers_sealed_client_when_server_has_shared_client( + self: "TestTemporaryMCPSessionEndpoints", + ) -> None: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_passthrough_authorization_code, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request: Final = MagicMock() + request.base_url = "https://litellm.example.com/" + request.headers = {} + server: Final = generate_mock_mcp_server_config_record(server_id="server-1").model_copy( + update={ + "auth_type": MCPAuth.true_passthrough, + "client_id": "shared-gateway-client", + "client_secret": "shared-secret", + "token_endpoint_auth_method": "client_secret_basic", + "token_url": "https://provider.example/oauth/token", + } + ) + admin_auth: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + fake_http_response: Final = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client: Final = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock(return_value=None) + mock_update_server: Final = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), + ), + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + ): + sealed: Final = seal_passthrough_authorization_code( + upstream_code="up-code", + client_id="ephemeral-client", + client_secret="ephemeral-secret", + mcp_server_id="server-1", + token_endpoint_auth_method="client_secret_post", + ) + result: Final = await mcp_token( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + grant_type="authorization_code", + code=sealed, + redirect_uri="https://example.com/callback", + client_id=None, + client_secret=None, + code_verifier="verifier", + refresh_token=None, + scope=None, + ) + + data: Final = fake_http_client.post.call_args.kwargs["data"] + headers: Final = fake_http_client.post.call_args.kwargs["headers"] + assert data["client_id"] == "ephemeral-client", data + assert data["client_secret"] == "ephemeral-secret", data + assert data["code"] == "up-code", data + assert data["redirect_uri"] == "https://litellm.example.com/callback", data + assert "Authorization" not in headers, headers + assert result.status_code == 200 + assert result.body == b'{"access_token":"tok","token_type":"Bearer"}' + @pytest.mark.asyncio async def test_mcp_token_refresh_grant_never_opens_sealed_code(self): """The minted client is unrecoverable outside the single authorization_code flow by @@ -3841,6 +3930,7 @@ class TestTemporaryMCPSessionEndpoints: refresh_token=None, scope=None, client_token_endpoint_auth_method=None, + sealed_client=False, ) @pytest.mark.asyncio @@ -3895,6 +3985,7 @@ class TestTemporaryMCPSessionEndpoints: refresh_token="rt-123", scope=None, client_token_endpoint_auth_method=None, + sealed_client=False, ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 21170abaeb1..ed89cc37544 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -33,6 +33,7 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException from litellm.proxy.utils import ProxyLogging from litellm.passthrough.main import raise_for_provider_status +from litellm.llms.bedrock.common_utils import BedrockError from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -2137,6 +2138,47 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) + + @pytest.mark.asyncio + async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = BedrockError( + status_code=400, + message="requestMetadata can contain a maximum of 16 items", + headers={"x-amzn-errortype": "ValidationException"}, + ) + + with pytest.raises(AwsErrorTypeProxyException) as exc_info: + await self._answer(error) + + assert exc_info.value.aws_error_type == "ValidationException" + assert "requestMetadata can contain a maximum of 16 items" in exc_info.value.error.message + assert exc_info.value.error.code == "400" + + @pytest.mark.asyncio + async def test_bedrock_error_without_an_aws_error_type_keeps_the_proxy_shape( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(BedrockError(status_code=400, message="bad request")) + + assert exc_info.value.code == "400" + + @pytest.mark.asyncio + async def test_error_with_a_non_ascii_header_value_keeps_the_proxy_shape( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = HTTPException(status_code=403, detail="model not allowed", headers={"x-model-group": "modèle"}) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"), diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 5fd088acd01..51c84618a74 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -13,7 +13,7 @@ import json import logging import os import re -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timedelta @@ -22,7 +22,9 @@ from types import MappingProxyType, SimpleNamespace from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +from prisma.errors import DataError from pydantic import JsonValue, TypeAdapter, ValidationError import litellm @@ -41,6 +43,7 @@ from litellm.proxy.proxy_server import ( validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, ) +from litellm.proxy.utils import evict_config_param from litellm.tracing.config import is_lens_tracing_enabled from .conftest import normalize @@ -1483,7 +1486,7 @@ def _install_fake_config_repo(monkeypatch, existing_row): captured["value"] = json.loads(data["update"]["param_value"]) class _FakeRepo: - def __init__(self, client): + def __init__(self, client, *, use_writer=False): self.table = _FakeTable() monkeypatch.setattr("litellm.proxy.proxy_server.ConfigRepository", _FakeRepo) @@ -1517,6 +1520,27 @@ async def test_ProxyConfig_save_environment_variables_merges_sets_and_deletes(mo assert "LITELLM_FAVICON_URL" not in written +@pytest.mark.asyncio +async def test_ProxyConfig_save_environment_variables_keeps_a_key_the_read_replica_has_not_caught_up_on(monkeypatch): + def config_table(row): + return SimpleNamespace(find_first=AsyncMock(return_value=SimpleNamespace(param_value=row)), upsert=AsyncMock()) + + writer_table = config_table({"SAVED_A_MOMENT_AGO": "ciphertext-new"}) + replica_table = config_table({}) + database = SimpleNamespace( + writer_db=SimpleNamespace(litellm_config=writer_table), db=SimpleNamespace(litellm_config=replica_table) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", database) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key") + + await ProxyConfig().save_environment_variables({"UI_LOGO_PATH": None}) + + written = replica_table.upsert.call_args.kwargs["data"]["update"]["param_value"] + assert json.loads(written) == {"SAVED_A_MOMENT_AGO": "ciphertext-new"} + + @pytest.mark.asyncio async def test_ProxyConfig_save_environment_variables_noop_without_db(monkeypatch): """With no DB configured the per-key write must do nothing (never touch the @@ -5078,6 +5102,418 @@ async def test_ProxyConfig_load_config_redacts_secret_litellm_setting_keeps_plai ) +_REVISION_TEST_KEYS: Final = ( + "REVISION_TEST_CONFIG_FILE", + "REVISION_TEST_DB_ONLY", + "REVISION_TEST_NULLED", + "REVISION_TEST_KEPT", + "revision_test_lower", + "REVISION_TEST_LOWER", + "GOOGLE_CLIENT_ID", + "PROXY_BASE_URL", +) + + +def _revision_test_environment() -> dict[str, str | None]: + return {key: os.environ.get(key) for key in _REVISION_TEST_KEYS} + + +@pytest.fixture +def revision_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-proxy-config-test-salt") + for key in _REVISION_TEST_KEYS: + monkeypatch.setenv(key, "restored-by-monkeypatch") + monkeypatch.delenv(key) + + +def _environment_variables_row( + values: Mapping[str, str | None] | None, revision: int | None +) -> SimpleNamespace | None: + if values is None: + return None + return SimpleNamespace( + param_name="environment_variables", + param_value={key: None if value is None else encrypt_value_helper(value) for key, value in values.items()}, + reload_revision=revision, + ) + + +_CONFIG_SECTIONS: Final = ("general_settings", "router_settings", "litellm_settings", "environment_variables") + + +async def _evict_config_rows() -> None: + for section in _CONFIG_SECTIONS: + await evict_config_param(section) + + +def _config_database( + environment_variables_read: Callable[[], Awaitable[SimpleNamespace | None]], + writer_find_unique: AsyncMock, + other_section_read: Callable[[], Awaitable[None]] | None = None, +) -> MagicMock: + async def get_generic_data(key: str, value: str, table_name: str) -> SimpleNamespace | None: + if value == "environment_variables": + return await environment_variables_read() + return None if other_section_read is None else await other_section_read() + + prisma_client: Final = MagicMock() + prisma_client.get_generic_data = get_generic_data + prisma_client.writer_db.litellm_config.find_unique = writer_find_unique + return prisma_client + + +def _rows_database(cached: SimpleNamespace | None, stored: SimpleNamespace | None) -> MagicMock: + async def read_cached() -> SimpleNamespace | None: + return cached + + return _config_database(read_cached, AsyncMock(return_value=stored)) + + +async def _sync_from(pc: ProxyConfig, prisma_client: MagicMock) -> None: + await _evict_config_rows() + await pc._update_config_from_db(prisma_client, {}, store_model_in_db=True) + await _evict_config_rows() + + +async def _sync_cached_and_stored_environment_variables_rows( + pc: ProxyConfig, cached: SimpleNamespace | None, stored: SimpleNamespace | None +) -> None: + await _sync_from(pc, _rows_database(cached, stored)) + + +async def _sync_environment_variables_row( + pc: ProxyConfig, values: Mapping[str, str | None] | None, revision: int | None +) -> None: + row: Final = _environment_variables_row(values, revision) + await _sync_cached_and_stored_environment_variables_rows(pc, cached=row, stored=row) + + +async def _init_non_llm_objects_from_an_empty_database(pc: ProxyConfig) -> None: + await _evict_config_rows() + await pc._init_non_llm_objects_in_db(prisma_client=_rows_database(None, None)) + await _evict_config_rows() + + +async def _sync_sso_settings_row(pc: ProxyConfig, values: Mapping[str, str] | None) -> None: + prisma_client: Final = MagicMock() + prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=None + if values is None + else SimpleNamespace(sso_settings={key: encrypt_value_helper(value) for key, value in values.items()}) + ) + await pc._init_sso_settings_in_db(prisma_client=prisma_client) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_newer_environment_variables_row_restores_the_keys_it_no_longer_stores(monkeypatch): + pc: Final = ProxyConfig() + monkeypatch.setenv("REVISION_TEST_CONFIG_FILE", "from-config-file") + + await _sync_environment_variables_row( + pc, + { + "REVISION_TEST_CONFIG_FILE": "from-db", + "REVISION_TEST_DB_ONLY": "db-only", + "REVISION_TEST_NULLED": "nulled-later", + "REVISION_TEST_KEPT": "kept-1", + "revision_test_lower": "lower", + }, + revision=1, + ) + assert _revision_test_environment() == { + "REVISION_TEST_CONFIG_FILE": "from-db", + "REVISION_TEST_DB_ONLY": "db-only", + "REVISION_TEST_NULLED": "nulled-later", + "REVISION_TEST_KEPT": "kept-1", + "revision_test_lower": "lower", + "REVISION_TEST_LOWER": "lower", + "GOOGLE_CLIENT_ID": None, + "PROXY_BASE_URL": None, + } + + await _sync_environment_variables_row( + pc, {"REVISION_TEST_NULLED": None, "REVISION_TEST_KEPT": "kept-2"}, revision=2 + ) + assert _revision_test_environment() == { + "REVISION_TEST_CONFIG_FILE": "from-config-file", + "REVISION_TEST_DB_ONLY": None, + "REVISION_TEST_NULLED": None, + "REVISION_TEST_KEPT": "kept-2", + "revision_test_lower": None, + "REVISION_TEST_LOWER": None, + "GOOGLE_CLIENT_ID": None, + "PROXY_BASE_URL": None, + } + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_cached_environment_variables_row_older_than_the_stored_one_changes_nothing(): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"REVISION_TEST_KEPT": "newer"}, revision=5) + + await _sync_cached_and_stored_environment_variables_rows( + pc, + cached=_environment_variables_row( + {"REVISION_TEST_KEPT": "older", "REVISION_TEST_DB_ONLY": "older-only"}, revision=4 + ), + stored=_environment_variables_row({"REVISION_TEST_KEPT": "newer"}, revision=5), + ) + + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_DB_ONLY")) == ("newer", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +@pytest.mark.parametrize("cached_revision", [None, 4]) +@pytest.mark.parametrize( + "writer_error", + [ + httpx.ConnectError("Can't reach database server"), + DataError( + {"user_facing_error": {"message": "Error querying the database: FATAL: sorry, too many clients already"}} + ), + ], + ids=["unreachable", "out_of_connections"], +) +async def test_a_writer_that_cannot_answer_keeps_the_applied_environment_variables_when_the_cached_row_is_behind( + cached_revision, writer_error +): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"REVISION_TEST_KEPT": "newer"}, revision=5) + cached: Final = _environment_variables_row( + {"REVISION_TEST_KEPT": "older", "REVISION_TEST_DB_ONLY": "older-only"}, revision=cached_revision + ) + + prisma_client: Final = _rows_database(cached, stored=None) + prisma_client.writer_db.litellm_config.find_unique = AsyncMock(side_effect=writer_error) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + await _sync_from(pc, prisma_client) + + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_DB_ONLY")) == ("newer", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_environment_variables_row_read_before_an_overlapping_newer_sync_is_not_applied_after_it(): + pc: Final = ProxyConfig() + older_values: Final = {"REVISION_TEST_KEPT": "older", "REVISION_TEST_DB_ONLY": "removed-by-the-newer-row"} + await _sync_environment_variables_row(pc, older_values, revision=5) + older_row: Final = _environment_variables_row(older_values, revision=5) + newer_row: Final = _environment_variables_row({"REVISION_TEST_KEPT": "newer"}, revision=6) + older_environment_row_read: Final = asyncio.Event() + newer_sync_finished: Final = asyncio.Event() + + async def read_older_row() -> SimpleNamespace | None: + older_environment_row_read.set() + return older_row + + async def wait_for_the_newer_sync() -> None: + await newer_sync_finished.wait() + + older_sync: Final = asyncio.create_task( + _sync_from( + pc, _config_database(read_older_row, AsyncMock(return_value=older_row), wait_for_the_newer_sync) + ) + ) + await older_environment_row_read.wait() + await _sync_from(pc, _rows_database(newer_row, newer_row)) + newer_sync_finished.set() + await older_sync + + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_DB_ONLY")) == ("newer", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_older_row_a_later_sync_read_before_a_newer_row_applied_is_not_applied_after_it(): + pc: Final = ProxyConfig() + older_values: Final = {"REVISION_TEST_KEPT": "older", "REVISION_TEST_DB_ONLY": "removed-by-the-newer-row"} + await _sync_environment_variables_row(pc, older_values, revision=5) + older_row: Final = _environment_variables_row(older_values, revision=5) + newer_row: Final = _environment_variables_row({"REVISION_TEST_KEPT": "newer"}, revision=6) + later_sync_read_the_older_row: Final = asyncio.Event() + earlier_sync_finished: Final = asyncio.Event() + + async def read_newer_row_after_the_later_sync_read() -> SimpleNamespace | None: + await later_sync_read_the_older_row.wait() + return newer_row + + async def read_older_row() -> SimpleNamespace | None: + later_sync_read_the_older_row.set() + return older_row + + async def wait_for_the_earlier_sync() -> None: + await earlier_sync_finished.wait() + + earlier_sync: Final = asyncio.create_task( + _sync_from(pc, _config_database(read_newer_row_after_the_later_sync_read, AsyncMock(return_value=newer_row))) + ) + await asyncio.sleep(0) + later_sync: Final = asyncio.create_task( + _sync_from(pc, _config_database(read_older_row, AsyncMock(return_value=newer_row), wait_for_the_earlier_sync)) + ) + await earlier_sync + earlier_sync_finished.set() + await later_sync + + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_DB_ONLY")) == ("newer", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_writer_row_read_while_an_overlapping_sync_applies_a_newer_row_is_not_applied_after_it(): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"REVISION_TEST_KEPT": "applied"}, revision=5) + writer_values: Final = {"REVISION_TEST_KEPT": "writer", "REVISION_TEST_DB_ONLY": "removed-by-the-newest-row"} + writer_row: Final = _environment_variables_row(writer_values, revision=6) + newest_row: Final = _environment_variables_row({"REVISION_TEST_KEPT": "newest"}, revision=7) + writer_read_started: Final = asyncio.Event() + earlier_sync_finished: Final = asyncio.Event() + + async def read_newest_row_once_the_writer_read_started() -> SimpleNamespace | None: + await writer_read_started.wait() + return newest_row + + async def read_cached_row_behind_the_applied_one() -> SimpleNamespace | None: + return _environment_variables_row(writer_values, revision=4) + + async def read_writer_row_finished_after_the_earlier_sync(where: Mapping[str, str]) -> SimpleNamespace | None: + writer_read_started.set() + await earlier_sync_finished.wait() + return writer_row + + earlier_sync: Final = asyncio.create_task( + _sync_from( + pc, _config_database(read_newest_row_once_the_writer_read_started, AsyncMock(return_value=newest_row)) + ) + ) + await asyncio.sleep(0) + later_sync: Final = asyncio.create_task( + _sync_from( + pc, + _config_database( + read_cached_row_behind_the_applied_one, + AsyncMock(side_effect=read_writer_row_finished_after_the_earlier_sync), + ), + ) + ) + await earlier_sync + earlier_sync_finished.set() + await later_sync + + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_DB_ONLY")) == ("newest", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_environment_variables_row_whose_stored_revision_went_backwards_is_applied(): + pc: Final = ProxyConfig() + await _sync_environment_variables_row( + pc, + {"REVISION_TEST_KEPT": "before-the-restore", "REVISION_TEST_DB_ONLY": "dropped-by-the-restore"}, + revision=500, + ) + + await _sync_environment_variables_row(pc, {"REVISION_TEST_KEPT": "restored"}, revision=1) + + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_DB_ONLY")) == ("restored", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_environment_variables_row_without_a_revision_only_sets_keys_a_later_revision_can_restore(): + pc: Final = ProxyConfig() + + await _sync_environment_variables_row( + pc, {"REVISION_TEST_KEPT": "kept", "REVISION_TEST_NULLED": "from-an-older-pod"}, revision=None + ) + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_NULLED")) == ( + "kept", + "from-an-older-pod", + ) + + await _sync_environment_variables_row(pc, {"REVISION_TEST_KEPT": "kept"}, revision=3) + assert (os.environ.get("REVISION_TEST_KEPT"), os.environ.get("REVISION_TEST_NULLED")) == ("kept", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_stale_environment_variables_row_without_a_revision_cannot_bring_back_a_removed_key(): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"REVISION_TEST_DB_ONLY": "db-only"}, revision=1) + await _sync_environment_variables_row(pc, {"REVISION_TEST_DB_ONLY": None}, revision=2) + + await _sync_cached_and_stored_environment_variables_rows( + pc, + cached=_environment_variables_row({"REVISION_TEST_DB_ONLY": "db-only"}, revision=None), + stored=_environment_variables_row({"REVISION_TEST_DB_ONLY": None}, revision=2), + ) + + assert os.environ.get("REVISION_TEST_DB_ONLY") is None + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_removed_key_does_not_come_back_with_a_value_set_after_the_first_sync_without_a_row(monkeypatch): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, None, revision=None) + + monkeypatch.setenv("REVISION_TEST_DB_ONLY", "set-at-runtime") + await _sync_environment_variables_row(pc, {"REVISION_TEST_DB_ONLY": "set-at-runtime"}, revision=1) + monkeypatch.delenv("REVISION_TEST_DB_ONLY") + await _sync_environment_variables_row(pc, {"REVISION_TEST_DB_ONLY": None}, revision=2) + + assert os.environ.get("REVISION_TEST_DB_ONLY") is None + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_sso_key_only_the_environment_variables_row_stored_is_unset_by_the_next_sso_sync(): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"PROXY_BASE_URL": "https://from-the-row"}, revision=1) + await _sync_sso_settings_row(pc, None) + + await _sync_environment_variables_row(pc, {"PROXY_BASE_URL": None}, revision=2) + await _sync_sso_settings_row(pc, None) + + assert os.environ.get("PROXY_BASE_URL") is None + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_sso_key_the_sso_settings_table_holds_stays_set_while_the_environment_variables_row_drops_it(): + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"GOOGLE_CLIENT_ID": "from-the-row"}, revision=1) + await _sync_sso_settings_row(pc, {"google_client_id": "from-the-row"}) + + await _sync_environment_variables_row(pc, {}, revision=2) + set_between_the_two_syncs: Final = os.environ.get("GOOGLE_CLIENT_ID") + await _sync_sso_settings_row(pc, {"google_client_id": "from-the-row"}) + + assert (set_between_the_two_syncs, os.environ.get("GOOGLE_CLIENT_ID")) == ("from-the-row", "from-the-row") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("revision_test_environment") +async def test_sso_key_the_environment_variables_row_drops_is_unset_when_sso_settings_are_not_loaded_from_the_db( + monkeypatch, +): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]}) + pc: Final = ProxyConfig() + await _sync_environment_variables_row(pc, {"GOOGLE_CLIENT_ID": "from-the-row"}, revision=1) + await _init_non_llm_objects_from_an_empty_database(pc) + + await _sync_environment_variables_row(pc, {}, revision=2) + await _init_non_llm_objects_from_an_empty_database(pc) + + assert os.environ.get("GOOGLE_CLIENT_ID") is None + + # --------------------------------------------------------------------------- # ProxyConfig agents from config.yaml # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index 2d86ff6a7c2..a92986392e9 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -51,6 +51,7 @@ def _install_litellm_config(mock_prisma: MagicMock) -> MagicMock: table.upsert = AsyncMock(return_value=None) table.delete = AsyncMock() mock_prisma.db.litellm_config = table + mock_prisma.writer_db.litellm_config = table return table diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 9e8910e44de..d5bacd9fd3c 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -437,15 +437,16 @@ def _decisions_provider_entry(provider: str) -> ProviderCreateInfo: return entry -def test_typesafe_provider_fields(): - typesafe: Final = _decisions_provider_entry("TypeSafe") +@pytest.mark.parametrize("display_name,provider", (("TypeSafe", LlmProviders.TYPESAFE), ("Nace", LlmProviders.NACE))) +def test_decisions_provider_fields(display_name: str, provider: LlmProviders) -> None: + entry: Final = _decisions_provider_entry(display_name) - assert typesafe.provider_display_name == "TypeSafe" - assert typesafe.litellm_provider == LlmProviders.TYPESAFE.value - assert typesafe.default_model_placeholder is not None - assert typesafe.default_model_placeholder.startswith("typesafe/") + assert entry.provider_display_name == display_name + assert entry.litellm_provider == provider.value + assert entry.default_model_placeholder is not None + assert entry.default_model_placeholder.startswith(f"{provider.value}/") - fields_by_key: Final = {f.key: f for f in typesafe.credential_fields} + fields_by_key: Final = {f.key: f for f in entry.credential_fields} assert fields_by_key["api_key"].required is True assert fields_by_key["api_key"].field_type == "password" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 8a01feb7492..03be269fdef 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -483,7 +483,12 @@ def _advertised_cost_split(headers: Mapping[str, str]) -> list[tuple[float, floa @pytest.mark.asyncio @pytest.mark.parametrize( "background, stored, status, expected_billed_events", - [(True, True, "completed", 1), (True, True, "in_progress", 0), (True, False, "completed", 0), (False, True, "completed", 0)], + [ + (True, True, "completed", 1), + (True, True, "in_progress", 0), + (True, False, "completed", 0), + (False, True, "completed", 0), + ], ) async def test_reading_a_finished_background_response_bills_it( monkeypatch: pytest.MonkeyPatch, background: bool, stored: bool, status: str, expected_billed_events: int @@ -534,8 +539,12 @@ async def test_reading_a_finished_background_response_bills_it( prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), ): await _bill_finished_background_response( @@ -665,9 +674,15 @@ async def _read_twice_on_two_pods( router.aget_responses = AsyncMock(side_effect=AssertionError("billing must not read the provider again")) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), ): await asyncio.gather( @@ -705,7 +720,9 @@ async def test_a_finished_response_the_create_could_not_store_is_billed_once_und monkeypatch, table, response_id, - UserAPIKeyAuth(api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42"), + UserAPIKeyAuth( + api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42" + ), ) (billed,) = payloads @@ -749,7 +766,9 @@ async def test_a_finished_read_never_bills_a_response_already_billed_or_billed_b ) table: Final = _UniqueManagedObjectTable(stored) - payloads: Final = await _read_twice_on_two_pods(monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader")) + payloads: Final = await _read_twice_on_two_pods( + monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader") + ) assert payloads == () assert table.rows == stored @@ -881,9 +900,15 @@ async def test_a_finished_read_or_cancel_bills_usage_only_the_configured_deploym "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("the configured deployment answers billing reads")), ): await getattr(endpoints, endpoint_name)( @@ -944,9 +969,15 @@ async def test_deleting_a_finished_unbilled_background_response_bills_it_before_ "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): result: Final = await delete_response( response_id=stored.response_id, @@ -992,7 +1023,9 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca caller: Final = UserAPIKeyAuth(api_key="sk-caller", user_role=caller_role) async def pipeline(**_: object) -> DeleteResponseResult: - await security.async_pre_call_hook(caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses") + await security.async_pre_call_hook( + caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses" + ) return DeleteResponseResult(id=stored.response_id, object="response", deleted=True) processor: Final = MagicMock() @@ -1005,9 +1038,15 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): with contextlib.suppress(HTTPException): await delete_response( @@ -1061,9 +1100,15 @@ def _capture_spend_events(monkeypatch: pytest.MonkeyPatch) -> _SpendEvents: @contextlib.contextmanager def _proxy_over(prisma_client: MagicMock, proxy_logging: MagicMock, router: MagicMock): with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): yield @@ -1214,7 +1259,9 @@ async def test_a_streamed_background_create_is_billed_wherever_it_sets_backgroun "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", return_value=False, ), - patch("litellm.proxy.proxy_server.select_data_generator", provider_events), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.select_data_generator", provider_events + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): await responses_api( request=_request_for("POST", "/v1/responses", body), @@ -1332,7 +1379,12 @@ async def test_background_create_the_provider_confirms_is_stored_under_its_provi } request = Request( - {"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"content-type", b"application/json")]}, + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"content-type", b"application/json")], + }, receive, ) @@ -1923,9 +1975,327 @@ class TestResponsesWSFirstFrameModelAuth: assert routed["input"] == replayed_input assert routed["previous_response_id"] == "resp_prev" assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini" + assert processor.common_processing_pre_call_logic.await_args.kwargs["skip_guardrails"] is True + assert callable(routed["responses_ws_frame_guard"]) assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket" ws.close.assert_not_awaited() + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrail_walk_on_the_frame_body_and_returns_its_records(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + assert call_type == "aresponses" and guardrails_only is True + assert "type" not in data and data["model"] == "gpt-4o-mini" + assert data["litellm_metadata"]["user_api_key_hash"] == "hash" + assert "standard_logging_guardrail_information" not in data["litellm_metadata"] + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] + return {**data, "input": data["input"] + " [guarded]"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + admitted = { + "model": "gpt-4o-mini", + "litellm_metadata": { + "user_api_key_hash": "hash", + "standard_logging_guardrail_information": [{"guardrail_name": "admission"}], + }, + } + guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj) + + first = await guard(json.dumps({"type": "response.create", "model": "alias", "input": "hi"})) + second = await guard(json.dumps({"type": "response.create", "input": "again"})) + + assert json.loads(first.frame) == {"type": "response.create", "model": "alias", "input": "hi [guarded]"} + assert first.guardrail_information == ({"guardrail_name": "turn"},) + assert json.loads(second.frame)["input"] == "again [guarded]" + assert admitted["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_name": "admission"} + ] + + @pytest.mark.asyncio + async def test_frame_guard_bills_a_rejected_frame_as_a_failed_request(self): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + raise ValueError("blocked") + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + proxy_logging_obj.post_call_failure_hook = AsyncMock() + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + rate_limiter = MagicMock() + rate_limiter.async_post_call_failure_hook = AsyncMock() + user_api_key_dict = MagicMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj + ) + + with ( + patch.object(litellm, "callbacks", [rate_limiter, spend_logger]), + pytest.raises(ValueError, match="blocked"), + ): + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + rate_limiter.async_post_call_failure_hook.assert_not_awaited() + spend_logger.async_post_call_failure_hook.assert_awaited_once() + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None}) + assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value + assert str(failure["original_exception"]) == "blocked" + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + assert failure["request_data"]["input"] == "hi" + assert len(failure["request_data"]["litellm_call_id"]) == 36 + + @pytest.mark.asyncio + async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + tools = [{"type": "function", "name": "a"}, {"type": "function", "name": "b"}] + + guarded = await guard(json.dumps({"type": "response.create", "response": {"input": "hi", "tools": tools}})) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["input"] == "hi" + assert json.loads(guarded.frame) == { + "type": "response.create", + "response": {"input": "hi [guarded]", "tools": tools[:1]}, + } + + @pytest.mark.asyncio + async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + metadata = { + "user_api_key_metadata": {"disable_global_guardrails": True}, + "disable_global_guardrails": True, + "headers": {"x-litellm-tags": "forged"}, + "tag": "kept", + } + + await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["metadata"] == {"tag": "kept"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "forbidden", "guardrails": []}, + {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, + ], + ) + async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + MagicMock(), + proxy_logging_obj, + ) + + await guard(json.dumps(frame)) + + scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"] + assert "guardrails" not in scanned and "policies" not in scanned + assert "guardrails" not in (scanned.get("metadata") or {}) + assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + + @pytest.mark.asyncio + async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): + import threading + + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + span = threading.RLock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"litellm_parent_otel_span": span}}, + MagicMock(), + proxy_logging_obj, + ) + + guarded = await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + assert ( + proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["litellm_metadata"]["litellm_parent_otel_span"] + is span + ) + + @pytest.mark.asyncio + async def test_frame_guard_rejects_a_turn_a_guardrail_rerouted_to_another_model(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + return {**data, "model": "onprem-llama"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match="rerouted this turn to 'onprem-llama'"), + ): + await guard(json.dumps({"type": "response.create", "input": "my SSN is 123-45-6789"})) + + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + assert "rerouted this turn to 'onprem-llama'" in str(failure["original_exception"]) + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_drops_a_client_s_guardrail_cost_records_so_a_blocked_frame_bills_the_real_scan( + self, nested: bool + ): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + raise ValueError("blocked") + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + body = { + "input": "forbidden", + "metadata": {"trace": "t1", "standard_logging_guardrail_information": [], "model_info": {"x": 1}}, + } + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"): + await guard(json.dumps(frame)) + + billed = spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"] + assert billed["metadata"] == {"trace": "t1"} + assert billed["litellm_metadata"]["standard_logging_guardrail_information"] == [{"guardrail_cost": 0.5}] + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_forwards_a_frame_without_the_instructions_a_guardrail_deleted(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {key: value for key, value in data.items() if key != "instructions"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + body = {"instructions": "SECRET system prompt", "input": "hi", "store": False} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with patch.object(litellm, "callbacks", []): + guarded = await guard(json.dumps(frame)) + + forwarded = json.loads(guarded.frame) + assert forwarded["type"] == "response.create" + sent_body = forwarded["response"] if nested else {k: v for k, v in forwarded.items() if k != "type"} + assert sent_body == {"input": "hi", "store": False} + + @pytest.mark.asyncio + async def test_frame_guard_hands_back_the_pii_tokens_presidio_masked_with_and_drops_forged_ones(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + assert data["metadata"] == {"trace": "t1"}, "a client cannot seed pii_tokens" + data["metadata"]["pii_tokens"] = {"": "alice@example.com"} + return {**data, "input": "email please"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "input": "email alice@example.com please", + "metadata": {"trace": "t1", "pii_tokens": {"": "mallory@example.com"}}, + } + ) + ) + + assert json.loads(guarded.frame)["input"] == "email please" + assert dict(guarded.pii_tokens) == {"": "alice@example.com"} + + @pytest.mark.asyncio + async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): + await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"): + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + spend_logger.async_post_call_failure_hook.assert_awaited_once() + @pytest.mark.asyncio @pytest.mark.parametrize("provider_rejected", [True, False]) async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool): @@ -3619,16 +3989,18 @@ class TestResponsesInputTokens: model_list=[ { "model_name": "gpt-5.2", - "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "additional_drop_params": ["tools"]}, + "litellm_params": { + "model": "openai/gpt-5.2", + "api_key": "sk-test", + "additional_drop_params": ["tools"], + }, }, {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test"}}, ], fallbacks=[{"gpt-5.2": ["backup"]}], ) - response, counter = self._post_input_tokens( - {"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router - ) + response, counter = self._post_input_tokens({"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router) assert response.status_code == 200, response.text assert counter.await_args.kwargs["request"].tools is None diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 1d197602543..679e1bfd85a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,7 +2,7 @@ import asyncio import copy import datetime import json -from collections.abc import Mapping +from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic +from litellm.types.utils import CachingDetails from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -1724,6 +1725,29 @@ class TestProxyBaseLLMRequestProcessing: assert "x-litellm-classifier-cost" not in headers + def test_get_custom_headers_drops_a_provider_header_http_cannot_carry(self): + user_api_key_dict: Final = MagicMock(spec=UserAPIKeyAuth) + user_api_key_dict.tpm_limit = None + user_api_key_dict.rpm_limit = None + user_api_key_dict.max_budget = None + user_api_key_dict.spend = 0 + + headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id="call-1", + **{ + "llm_provider-x-amzn-requestid": "req-1", + "llm_provider-x-provider-region": "東京リージョン", + "llm_provider-x-provider-city": "Zürich", + }, + ) + response: Final = Response() + response.headers.update(headers) + + assert response.headers["llm_provider-x-amzn-requestid"] == "req-1" + assert response.headers["llm_provider-x-provider-city"] == "Zürich" + assert "llm_provider-x-provider-region" not in response.headers + def test_get_cost_breakdown_from_logging_obj_helper(self): """ Test the helper function that extracts cost breakdown information. @@ -4600,6 +4624,32 @@ class TestStreamCloseOnDisconnect: timeout=5, ) + async def test_cancelling_the_first_chunk_wait_stops_the_upstream_read(self): + reads: list[int] = [] + closed = asyncio.Event() + + async def slow_upstream(): + try: + while True: + reads.append(len(reads)) + await asyncio.sleep(0.01) + yield "data: never\n\n" + finally: + closed.set() + + waiting = asyncio.ensure_future( + _buffer_first_chunk_honoring_disconnect(slow_upstream(), request=self._request_that_stays_connected()) + ) + await asyncio.sleep(0.05) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + reads_at_cancel = len(reads) + await asyncio.sleep(0.05) + + assert closed.is_set() + assert len(reads) == reads_at_cancel + async def test_buffer_first_chunk_raises_sentinel_and_closes_on_disconnect(self): closed = asyncio.Event() @@ -6009,7 +6059,7 @@ class TestAllmPassthroughStreamingProviderGate: } return ProxyBaseLLMRequestProcessing(data=data) - async def _run(self, processing_obj, monkeypatch, chunks, stream=None): + async def _run(self, processing_obj, monkeypatch, chunks, stream=None, proxy_logging_obj=None, request=None): import litellm.proxy.common_request_processing as crp from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth @@ -6027,14 +6077,14 @@ class TestAllmPassthroughStreamingProviderGate: monkeypatch.setattr(crp, "route_request", fake_route_request) - proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj = proxy_logging_obj if proxy_logging_obj is not None else MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) proxy_logging_obj.update_request_status = AsyncMock(return_value=None) proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) proxy_logging_obj.post_call_success_hook = AsyncMock() return await processing_obj.base_process_llm_request( - request=MagicMock(spec=Request, headers={}), + request=request if request is not None else MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"), route_type="allm_passthrough_route", @@ -6100,15 +6150,144 @@ class TestAllmPassthroughStreamingProviderGate: streamed = [chunk async for chunk in result.body_iterator] assert streamed == chunks - @pytest.mark.asyncio - async def test_bedrock_converse_stream_is_buffered_through_handler( - self, monkeypatch - ): - processing_obj = self._build_processing_obj( - "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" - ) - chunks = [b"raw-1", b"raw-2"] + @staticmethod + def _converse_frame(event_type: str, payload: dict) -> bytes: + import struct + from binascii import crc32 + headers = b"".join( + bytes((len(name),)) + name + b"\x07" + struct.pack("!H", len(value)) + value + for name, value in ( + (b":event-type", event_type.encode()), + (b":content-type", b"application/json"), + (b":message-type", b"event"), + ) + ) + body = json.dumps(payload).encode() + prelude = struct.pack("!II", 16 + len(headers) + len(body), len(headers)) + message = prelude + struct.pack("!I", crc32(prelude)) + headers + body + return message + struct.pack("!I", crc32(message)) + + def _converse_stream(self) -> list[bytes]: + return [ + self._converse_frame("messageStart", {"role": "assistant"}), + self._converse_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "Hello"}}), + self._converse_frame("contentBlockStop", {"contentBlockIndex": 0}), + self._converse_frame("messageStop", {"stopReason": "end_turn"}), + ] + + @staticmethod + def _apply_guardrail(**kwargs): + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _ApplyGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _ApplyGuardrail(guardrail_name="scan", event_hook="post_call", default_on=True, **kwargs) + + @staticmethod + def _success_hook_guardrail(): + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _SuccessHookGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + return _SuccessHookGuardrail(guardrail_name="whole", event_hook="post_call", default_on=True) + + @staticmethod + def _native_hooks_guardrail(): + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _NativeHooksGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _NativeHooksGuardrail(guardrail_name="native", event_hook="post_call", default_on=True) + + @classmethod + def _release_on_scan_guardrail(cls): + guardrail = cls._apply_guardrail() + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = True + return guardrail + + @classmethod + def _parallel_guardrail(cls): + return cls._apply_guardrail(run_in_parallel=True) + + @staticmethod + def _holding_release_on_scan_guardrail(): + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _HoldingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for item in response: + yield item + + guardrail = _HoldingGuardrail(guardrail_name="held", event_hook="post_call", default_on=True) + guardrail.streaming_buffer_release_on_scan = True + return guardrail + + @staticmethod + def _recording_upstream(frames: list[bytes], events: list[str]): + class _Upstream: + def __init__(self) -> None: + self._frames = iter(frames) + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + frame = next(self._frames, None) + if frame is None: + raise StopAsyncIteration + return frame + + async def aclose(self) -> None: + events.append("upstream closed") + + async def wait_for_spend_flush(self) -> bool: + events.append("spend logged") + return True + + return _Upstream() + + async def _run_converse_stream( + self, + monkeypatch, + guardrail, + chunks, + event_stream_hook=None, + stream=None, + failure_hook=None, + read_body=True, + request=None, + endpoint="model/us.amazon.nova-lite-v1:0/converse-stream", + model="us.amazon.nova-lite-v1:0", + post_call_guarded=True, + request_body=None, + ): + processing_obj = self._build_processing_obj("bedrock", endpoint) + processing_obj.data["litellm_logging_obj"].model = model + processing_obj.data["data"] = request_body + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + recorded: dict = {} + + async def relay_frames(response, user_api_key_dict, request_data, guardrails): + recorded["guardrails"] = list(guardrails) + async for item in response: + yield item + + self._proxy_logging_obj = MagicMock(spec=ProxyLogging) + self._proxy_logging_obj.async_post_call_event_stream_iterator_hook = event_stream_hook or relay_frames + self._proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None, side_effect=failure_hook) with patch.object( ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", @@ -6116,19 +6295,483 @@ class TestAllmPassthroughStreamingProviderGate: ), patch.object( ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", - return_value=True, + return_value=post_call_guarded, ), patch( "litellm.llms.bedrock.passthrough.guardrail_translation.handler." "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", new=AsyncMock(return_value=b"modified-body"), ) as mock_handler: - result = await self._run(processing_obj, monkeypatch, chunks) + result = await self._run( + processing_obj, + monkeypatch, + chunks, + stream=stream, + proxy_logging_obj=self._proxy_logging_obj, + request=request, + ) + body = ( + b"".join([chunk async for chunk in result.body_iterator]) + if isinstance(result, StreamingResponse) and read_body + else None + if isinstance(result, StreamingResponse) + else result.body + ) + return result, body, mock_handler, recorded + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "guardrail_factory", + [ + "_success_hook_guardrail", + "_native_hooks_guardrail", + "_mask_response_guardrail", + "_release_on_scan_guardrail", + "_holding_release_on_scan_guardrail", + "_parallel_guardrail", + ], + ) + @pytest.mark.parametrize( + ("endpoint", "model", "request_body"), + [ + ("model/us.amazon.nova-lite-v1:0/converse-stream", "us.amazon.nova-lite-v1:0", None), + ( + "model/us.anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream", + "us.anthropic.claude-sonnet-5-v1:0", + None, + ), + ( + "model/abc/invoke-with-response-stream", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc", + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 8, "messages": []}, + ), + ], + ids=["converse-stream", "claude-invoke-stream", "claude-invoke-stream-behind-an-application-inference-profile"], + ) + async def test_bedrock_stream_is_buffered_for_a_guardrail_that_needs_the_whole_response( + self, + monkeypatch: pytest.MonkeyPatch, + guardrail_factory: str, + endpoint: str, + model: str, + request_body: Mapping[str, object] | None, + ): + guardrail = ( + self._apply_guardrail(mask_response_content=True) + if guardrail_factory == "_mask_response_guardrail" + else getattr(self, guardrail_factory)() + ) + + result, body, mock_handler, recorded = await self._run_converse_stream( + monkeypatch, + guardrail, + [b"raw-1", b"raw-2"], + endpoint=endpoint, + model=model, + request_body=request_body, + ) assert isinstance(result, Response) assert not isinstance(result, StreamingResponse) - assert result.body == b"modified-body" + assert body == b"modified-body" assert result.headers["content-type"] == "application/vnd.amazon.eventstream" mock_handler.assert_awaited_once() + assert recorded == {} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "streaming_flags", + [ + {}, + {"streaming_buffer_release_on_scan": True}, + { + "streaming_buffer_until_moderated": True, + "streaming_buffer_release_on_scan": True, + "streaming_end_of_stream_only": True, + }, + ], + ids=["sampled", "release-on-scan-while-sampled", "release-on-scan-at-end-of-stream-only"], + ) + async def test_bedrock_converse_stream_relays_frames_for_an_apply_guardrail_guardrail( + self, monkeypatch, streaming_flags + ): + guardrail = self._apply_guardrail() + for name, value in streaming_flags.items(): + setattr(guardrail, name, value) + stream = b"".join(self._converse_stream()) + + result, body, mock_handler, recorded = await self._run_converse_stream( + monkeypatch, guardrail, [stream[:20], stream[20:]] + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert result.headers["content-type"] == "application/vnd.amazon.eventstream" + assert body == stream + assert recorded == {"guardrails": [(guardrail, False)]} + mock_handler.assert_not_awaited() + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_failure_after_the_first_frame_ends_with_an_exception_frame( + self, monkeypatch + ): + from botocore.eventstream import EventStreamBuffer + + from litellm.exceptions import GuardrailRaisedException + + async def block_after_first_frame(response, user_api_key_dict, request_data, guardrails): + async for item in response: + yield item + raise GuardrailRaisedException( + guardrail_name="scan", message="Violated guardrail policy", should_wrap_with_default_message=False + ) + + frames = self._converse_stream() + result, body, _, _ = await self._run_converse_stream( + monkeypatch, self._apply_guardrail(), frames, event_stream_hook=block_after_first_frame + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert body.startswith(frames[0]) + buffer = EventStreamBuffer() + buffer.add_data(body[len(frames[0]) :]) + messages = list(buffer) + assert [(m.headers[":message-type"], m.headers[":exception-type"]) for m in messages] == [ + ("exception", "validationException") + ] + assert json.loads(messages[0].payload) == {"message": "Violated guardrail policy"} + self._proxy_logging_obj.post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_failure_after_the_first_frame_logs_the_upstream_spend_first( + self, monkeypatch + ): + from litellm.exceptions import GuardrailRaisedException + + events: list[str] = [] + frames = self._converse_stream() + + async def block_after_first_frame(response, user_api_key_dict, request_data, guardrails): + async for item in response: + yield item + raise GuardrailRaisedException( + guardrail_name="scan", message="Violated guardrail policy", should_wrap_with_default_message=False + ) + + async def failure_hook(**kwargs): + events.append("failure hook") + + result, _, _, _ = await self._run_converse_stream( + monkeypatch, + self._apply_guardrail(), + frames, + event_stream_hook=block_after_first_frame, + stream=self._recording_upstream(frames, events), + failure_hook=failure_hook, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert events[:3] == ["upstream closed", "spend logged", "failure hook"] + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_block_raised_after_the_first_frame_runs_the_failure_hook(self, monkeypatch): + from botocore.eventstream import EventStreamBuffer + + async def block_after_first_frame(response, user_api_key_dict, request_data, guardrails): + async for item in response: + yield item + raise HTTPException(status_code=400, detail="Violated guardrail policy") + + frames = self._converse_stream() + result, body, _, _ = await self._run_converse_stream( + monkeypatch, self._apply_guardrail(), frames, event_stream_hook=block_after_first_frame + ) + + assert body.startswith(frames[0]) + buffer = EventStreamBuffer() + buffer.add_data(body[len(frames[0]) :]) + messages = list(buffer) + assert [(m.headers[":message-type"], m.headers[":exception-type"]) for m in messages] == [ + ("exception", "validationException") + ] + assert json.loads(messages[0].payload) == {"message": "Violated guardrail policy"} + self._proxy_logging_obj.post_call_failure_hook.assert_awaited_once() + assert self._proxy_logging_obj.post_call_failure_hook.await_args.kwargs["original_exception"].status_code == 400 + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_failing_after_an_empty_upstream_raises(self, monkeypatch): + async def fail_at_the_end(response, user_api_key_dict, request_data, guardrails): + async for item in response: + yield item + raise HTTPException(status_code=503, detail="guardrail unavailable") + + with pytest.raises(HTTPException) as raised: + await self._run_converse_stream(monkeypatch, self._apply_guardrail(), [], event_stream_hook=fail_at_the_end) + + assert raised.value.status_code == 503 + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_closed_before_its_body_starts_closes_the_guardrails_and_the_upstream( + self, monkeypatch + ): + events: list[str] = [] + frames = self._converse_stream() + + async def scan(response, user_api_key_dict, request_data, guardrails): + try: + async for item in response: + yield item + finally: + events.append("guardrails closed") + + result, _, _, _ = await self._run_converse_stream( + monkeypatch, + self._apply_guardrail(), + frames, + event_stream_hook=scan, + stream=self._recording_upstream(frames, events), + read_body=False, + ) + await result.upstream_generator.aclose() + + assert events == ["guardrails closed", "upstream closed"] + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_client_leaving_a_held_stream_before_its_first_frame_gets_a_499( + self, monkeypatch + ): + events: list[str] = [] + frames = self._converse_stream() + first_frame_read = asyncio.Event() + release_budget = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation_on_cancel", release_budget + ) + + class _SlowUpstream: + def __init__(self) -> None: + self._frames = iter(frames) + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + frame = next(self._frames, None) + if frame is None: + raise StopAsyncIteration + if first_frame_read.is_set(): + await asyncio.Event().wait() + first_frame_read.set() + return frame + + async def aclose(self) -> None: + events.append("upstream closed") + + async def wait_for_spend_flush(self) -> bool: + events.append("spend logged") + return True + + async def hold_until_the_end(response, user_api_key_dict, request_data, guardrails): + held = [item async for item in response] + for item in held: + yield item + + async def client_leaves_once_the_upstream_answers(): + await first_frame_read.wait() + return {"type": "http.disconnect"} + + request = MagicMock(spec=Request, headers={}) + request.receive = client_leaves_once_the_upstream_answers + + result, body, _, _ = await asyncio.wait_for( + self._run_converse_stream( + monkeypatch, + self._apply_guardrail(), + frames, + event_stream_hook=hold_until_the_end, + stream=_SlowUpstream(), + request=request, + ), + timeout=5, + ) + + assert isinstance(result, JSONResponse) + assert result.status_code == 499 + assert json.loads(body)["error"]["type"] == "client_disconnect" + assert events[-1] == "spend logged" + assert "upstream closed" in events + self._proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + release_budget.assert_not_awaited() + self._proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_not_awaited() + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_client_leaving_before_bedrock_answers_releases_what_the_request_holds( + self, monkeypatch + ): + events: list[str] = [] + release_budget = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation_on_cancel", release_budget + ) + + class _SilentUpstream: + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + await asyncio.Event().wait() + raise StopAsyncIteration + + async def aclose(self) -> None: + events.append("upstream closed") + + async def wait_for_spend_flush(self) -> bool: + return False + + async def client_leaves(): + return {"type": "http.disconnect"} + + request = MagicMock(spec=Request, headers={}) + request.receive = client_leaves + + result, _, _, _ = await asyncio.wait_for( + self._run_converse_stream( + monkeypatch, self._apply_guardrail(), [], stream=_SilentUpstream(), request=request + ), + timeout=5, + ) + + assert isinstance(result, JSONResponse) + assert result.status_code == 499 + assert "upstream closed" in events + release_budget.assert_awaited_once_with(None) + self._proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + + @pytest.mark.asyncio + async def test_bedrock_converse_stream_failure_before_the_first_frame_raises(self, monkeypatch): + events: list[str] = [] + frames = self._converse_stream() + + async def block_before_any_frame(response, user_api_key_dict, request_data, guardrails): + async for _ in response: + raise HTTPException(status_code=400, detail="Violated guardrail policy") + yield b"" + + with pytest.raises(HTTPException) as raised: + await self._run_converse_stream( + monkeypatch, + self._apply_guardrail(), + frames, + event_stream_hook=block_before_any_frame, + stream=self._recording_upstream(frames, events), + ) + + assert raised.value.status_code == 400 + assert events[-1] == "spend logged" + assert "upstream closed" in events + self._proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + + def _invoke_stream(self) -> list[bytes]: + import base64 + + return [ + self._converse_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + for event in ( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}, + ) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "request_body"), + [ + ("us.anthropic.claude-sonnet-5-v1:0", None), + ( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc", + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 8, "messages": []}, + ), + ], + ids=["claude-model-id", "application-inference-profile-with-a-claude-body"], + ) + async def test_bedrock_invoke_stream_of_a_claude_model_is_relayed_and_scanned_for_its_text( + self, monkeypatch: pytest.MonkeyPatch, model: str, request_body: Mapping[str, object] | None + ): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + guardrail = self._apply_guardrail() + scanned: dict[str, object] = {} + + async def scan_relayed( + response: AsyncIterable[object], + user_api_key_dict: ProxyUserAPIKeyAuth, + request_data: dict[str, object], + guardrails: Sequence[tuple[CustomGuardrail, bool]], + ) -> AsyncGenerator[object, None]: + items = [item async for item in response] + scanned["guardrails"] = list(guardrails) + scanned["texts"] = LlmPassthroughRouteHandler().get_streaming_scan_key(items).texts + for item in items: + yield item + + stream = b"".join(self._invoke_stream()) + result, body, mock_handler, _ = await self._run_converse_stream( + monkeypatch, + guardrail, + [stream[:20], stream[20:]], + event_stream_hook=scan_relayed, + endpoint=f"model/{model}/invoke-with-response-stream", + model=model, + request_body=request_body, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert result.headers["content-type"] == "application/vnd.amazon.eventstream" + assert body == stream + assert scanned == {"guardrails": [(guardrail, False)], "texts": ("Hello",)} + mock_handler.assert_not_awaited() + + @pytest.mark.asyncio + async def test_bedrock_invoke_stream_of_a_claude_model_without_a_post_call_guardrail_streams_as_received( + self, monkeypatch: pytest.MonkeyPatch + ): + result, body, mock_handler, recorded = await self._run_converse_stream( + monkeypatch, + self._apply_guardrail(), + [b"raw-1", b"raw-2"], + endpoint="model/us.anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream", + model="us.anthropic.claude-sonnet-5-v1:0", + post_call_guarded=False, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert body == b"raw-1raw-2" + assert recorded == {} + mock_handler.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "guardrail_factory"), + [ + ("meta.llama4-maverick-17b-instruct-v1:0", "_apply_guardrail"), + ("meta.llama4-maverick-17b-instruct-v1:0", "_success_hook_guardrail"), + ], + ids=["guardrail-the-relay-scans", "guardrail-that-needs-the-whole-response"], + ) + async def test_bedrock_invoke_stream_of_a_model_without_a_stream_translation_streams_as_received( + self, monkeypatch: pytest.MonkeyPatch, model: str, guardrail_factory: str + ): + result, body, mock_handler, recorded = await self._run_converse_stream( + monkeypatch, + getattr(self, guardrail_factory)(), + [b"raw-1", b"raw-2"], + endpoint=f"model/{model}/invoke-with-response-stream", + model=model, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert body == b"raw-1raw-2" + assert recorded == {} + mock_handler.assert_not_awaited() @pytest.mark.asyncio async def test_bedrock_invoke_stream_is_not_buffered(self, monkeypatch): @@ -10995,3 +11638,98 @@ class TestResponseCostFromLoggingObjForBridgedMessages: assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(expected_output) assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) + + +class TestResponseCostFromLoggingObjForCacheHits: + BODY: Final = MappingProxyType( + { + "id": "msg_cached", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "cached answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1200, "output_tokens": 400}, + } + ) + + @staticmethod + def _logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = Logging( + model="claude-haiku-5-5", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id="cache-hit-header", + ) + logging_obj.optional_params = {} + return logging_obj + + def test_a_v1_messages_cache_hit_costs_nothing_though_its_usage_prices_as_a_fresh_call(self): + fresh_cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(self.BODY), logging_obj=self._logging_obj() + ) + hit: Final = self._logging_obj() + hit.caching_details = {"cache_hit": True, "cache_duration_ms": 2.0} + + hit_cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(self.BODY), logging_obj=hit + ) + + assert fresh_cost is not None and fresh_cost > 0 + assert hit_cost == 0.0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("hidden_params", "caching_details", "expected_header"), + ( + ( + {"cache_hit": True, "response_cost": 0.0}, + CachingDetails(cache_hit=True, cache_duration_ms=2.0), + "0.0", + ), + ({"response_cost": 0.0}, None, None), + ), + ) + async def test_only_a_cache_hit_sends_zero_cost_headers( + self, + hidden_params: dict[str, object], + caching_details: CachingDetails | None, + expected_header: str | None, + ): + response: Final = litellm.ModelResponse() + response._hidden_params = hidden_params + logging_obj: Final = self._logging_obj() + logging_obj.caching_details = caching_details + + headers: Final = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data={"model": "gpt-5.4-mini"}, + request=Request({"type": "http", "method": "POST", "headers": []}), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + logging_obj=logging_obj, + version=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + cost_headers: Final = ( + "x-litellm-response-cost", + "x-litellm-response-cost-original", + "x-litellm-response-cost-input", + "x-litellm-response-cost-output", + ) + assert {name: headers.get(name) for name in cost_headers} == dict.fromkeys(cost_headers, expected_header) + + def test_a_fresh_zero_cost_v1_messages_call_still_sends_no_cost_header(self): + logging_obj: Final = self._logging_obj() + logging_obj.model_call_details["response_cost"] = 0.0 + + cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(self.BODY), logging_obj=logging_obj + ) + + assert cost == "" diff --git a/tests/unit/proxy/test_component_allowlists.py b/tests/unit/proxy/test_component_allowlists.py index 42da00fcb6a..eb058d01770 100644 --- a/tests/unit/proxy/test_component_allowlists.py +++ b/tests/unit/proxy/test_component_allowlists.py @@ -57,7 +57,7 @@ _PRE_EXISTING_ENV = {key: os.environ.get(key) for key in _THROWAWAY_ENV} for _key, _value in _THROWAWAY_ENV.items(): os.environ.setdefault(_key, _value) -from prometheus_client import make_asgi_app +from prometheus_client import CollectorRegistry, Counter, make_asgi_app # gateway/ and backend/ live at the repo root, not inside litellm/. _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) @@ -71,6 +71,7 @@ from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app +from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_app from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter for _key, _previous in _PRE_EXISTING_ENV.items(): @@ -355,6 +356,26 @@ def test_gateway_trim_keeps_metrics_mount(): "/metrics Mount must survive the gateway route trim" +@pytest.mark.parametrize("enabled", (False, True)) +def test_backend_serves_metrics_after_startup_when_enabled(enabled: bool) -> None: + registry: Final = CollectorRegistry() + failures: Final = Counter("litellm_postgres_failed_requests", "Failed database requests", registry=registry) + failures.inc(2) + + @asynccontextmanager + async def lifespan(application: Starlette) -> AsyncGenerator[None, None]: + if enabled: + application.mount("/metrics", make_metrics_asgi_app(registry)) + yield + + application: Final = FastAPI(lifespan=partial(_backend_lifespan, lifespan=lifespan)) + with TestClient(application) as client: + response: Final = client.get("/metrics") + assert response.status_code == (200 if enabled else 404), response.text + if enabled: + assert "litellm_postgres_failed_requests_total 2.0\n" in response.text + + def test_gateway_drops_ui_and_swagger_mounts(): """UI static and swagger Mounts must still be trimmed from the gateway.""" for path in ("/ui", "/_next", "/litellm-asset-prefix/_next", "/swagger"): diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 9566b6e30a2..89a35923be0 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -369,16 +369,16 @@ class TestProxyInitializationHelpers: assert config_file.resolve() in yielded_paths def test_dev_env_hot_reload_enabled_reads_flag(self, monkeypatch): - import litellm + from litellm._process_defaults import _dev_env_hot_reload_enabled monkeypatch.setenv("LITELLM_DEV_ENV_HOT_RELOAD", "True") - assert litellm._dev_env_hot_reload_enabled() is True + assert _dev_env_hot_reload_enabled() is True monkeypatch.setenv("LITELLM_DEV_ENV_HOT_RELOAD", "false") - assert litellm._dev_env_hot_reload_enabled() is False + assert _dev_env_hot_reload_enabled() is False monkeypatch.delenv("LITELLM_DEV_ENV_HOT_RELOAD", raising=False) - assert litellm._dev_env_hot_reload_enabled() is False + assert _dev_env_hot_reload_enabled() is False @patch("asyncio.run") @patch("builtins.print") @@ -526,6 +526,12 @@ class TestProxyInitializationHelpers: with patch("sys.platform", "linux"): assert ProxyInitializationHelpers._get_loop_type() == "uvloop" + @pytest.mark.parametrize(("env_value", "expected"), [("asyncio", "asyncio"), ("uvloop", "uvloop"), ("", "uvloop")]) + def test_get_loop_type_honors_litellm_event_loop_env(self, env_value: str, expected: str, monkeypatch) -> None: + monkeypatch.setattr("sys.platform", "linux") + monkeypatch.setenv("LITELLM_EVENT_LOOP", env_value) + assert ProxyInitializationHelpers._get_loop_type() == expected + def test_selected_loop_factory_imports_on_this_interpreter(self): loop_type = ProxyInitializationHelpers._get_loop_type() if loop_type is None: diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 80bc3ad1a34..2f2563b9ee4 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -3205,6 +3205,7 @@ async def test_update_config_success_callback_normalization(monkeypatch): self.db.litellm_config = MagicMock() self.db.litellm_config.find_first = AsyncMock(side_effect=fake_find_first) self.db.litellm_config.upsert = AsyncMock(side_effect=fake_upsert) + self.writer_db = self.db monkeypatch.setattr(proxy_server, "prisma_client", MockPrisma()) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index e824c5017c6..fbdd8d4d7fb 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -6691,10 +6691,9 @@ def test_settings_store_uppercases_db_env_vars(monkeypatch): monkeypatch.delenv(key, raising=False) proxy_config = ProxyConfig() - db_values = proxy_config._prepared_db_settings_values( - "environment_variables", {"dd_api_key": "test-api-key", "dd_site": "us5.datadoghq.com"} + proxy_config._apply_db_environment_variables( + {"dd_api_key": "test-api-key", "dd_site": "us5.datadoghq.com"}, revision=None, environment_before_db={}, read=1 ) - proxy_config.environment_variables.apply_db_row("environment_variables", db_values) env_vars = proxy_config.environment_variables.resolved() assert env_vars["DD_API_KEY"] == "test-api-key" @@ -9924,10 +9923,22 @@ class _FakeLitellmConfig: self.upsert_calls.append((name, value)) +class _LaggingReplicaLitellmConfig(_FakeLitellmConfig): + def __init__(self, replica_rows, writer): + super().__init__(initial_rows=replica_rows) + self.upsert = writer.upsert + + class _FakePrismaClient: - def __init__(self, initial_rows=None): + def __init__(self, initial_rows=None, replica_rows=None): + self.writer_db = mock.MagicMock() + self.writer_db.litellm_config = _FakeLitellmConfig(initial_rows=initial_rows) self.db = mock.MagicMock() - self.db.litellm_config = _FakeLitellmConfig(initial_rows=initial_rows) + self.db.litellm_config = ( + self.writer_db.litellm_config + if replica_rows is None + else _LaggingReplicaLitellmConfig(replica_rows, writer=self.writer_db.litellm_config) + ) self.jsonify_object = lambda obj: obj @@ -9937,8 +9948,8 @@ def _update_config_setup(monkeypatch): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth as auth_dep - def _install(initial_rows=None, store_model_in_db=True): - prisma = _FakePrismaClient(initial_rows=initial_rows) + def _install(initial_rows=None, store_model_in_db=True, replica_rows=None): + prisma = _FakePrismaClient(initial_rows=initial_rows, replica_rows=replica_rows) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", store_model_in_db) monkeypatch.setattr( @@ -9992,6 +10003,22 @@ def test_update_config_writes_only_sent_section(_update_config_setup): restore() +def test_update_config_keeps_an_environment_variable_the_read_replica_has_not_caught_up_on(_update_config_setup): + client, prisma, restore = _update_config_setup( + initial_rows={"environment_variables": {"SAVED_A_MOMENT_AGO": "enc:new"}}, + replica_rows={"environment_variables": {}}, + ) + try: + resp = client.post("/config/update", json={"environment_variables": {"ANOTHER_KEY": "value"}}) + assert resp.status_code == 200 + assert prisma.writer_db.litellm_config.rows["environment_variables"] == { + "SAVED_A_MOMENT_AGO": "enc:new", + "ANOTHER_KEY": "enc:value", + } + finally: + restore() + + def test_update_config_rejects_overlapping_routing_groups_before_writing(_update_config_setup): existing_groups = [{"group_name": "g1", "models": ["m1"], "routing_strategy": "least-busy"}] client, prisma, restore = _update_config_setup( diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index da799a307a6..512f58380ea 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -646,7 +646,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_config.update = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) @@ -731,7 +731,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_config.update = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) @@ -863,7 +863,7 @@ class TestProxySettingEndpoints: "PROXY_BASE_URL": "old_proxy_url", } ) - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( return_value=env_var_entry ) mock_prisma.db.litellm_config.update = AsyncMock() @@ -942,7 +942,7 @@ class TestProxySettingEndpoints: "PROXY_BASE_URL": "old_proxy_url", } ) - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( return_value=env_var_entry ) mock_prisma.db.litellm_config.update = AsyncMock() @@ -1013,7 +1013,7 @@ class TestProxySettingEndpoints: "MICROSOFT_CLIENT_SECRET": "test_existing_microsoft_secret", } ) - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( return_value=env_var_entry ) mock_prisma.db.litellm_config.update = AsyncMock() @@ -1083,7 +1083,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_config.update = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) @@ -1944,7 +1944,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = upsert_mock mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_config.update = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) @@ -2033,7 +2033,7 @@ class TestProxySettingEndpoints: } ) mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( return_value=env_var_entry ) mock_prisma.db.litellm_config.update = AsyncMock() @@ -2052,7 +2052,7 @@ class TestProxySettingEndpoints: ) assert response.status_code == 200 - mock_prisma.db.litellm_config.find_unique.assert_called_once() + mock_prisma.writer_db.litellm_config.find_unique.assert_called_once() mock_prisma.db.litellm_config.update.assert_called_once() update_call = mock_prisma.db.litellm_config.update.call_args updated_env_vars = json.loads(update_call.kwargs["data"]["param_value"]) @@ -2082,7 +2082,7 @@ class TestProxySettingEndpoints: "ANOTHER_ENV": "also_keep", } mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( return_value=env_var_entry ) mock_prisma.db.litellm_config.update = AsyncMock() @@ -2101,12 +2101,44 @@ class TestProxySettingEndpoints: ) assert response.status_code == 200 - mock_prisma.db.litellm_config.find_unique.assert_called_once() + mock_prisma.writer_db.litellm_config.find_unique.assert_called_once() mock_prisma.db.litellm_config.update.assert_called_once() update_call = mock_prisma.db.litellm_config.update.call_args updated_env_vars = json.loads(update_call.kwargs["data"]["param_value"]) assert updated_env_vars == env_var_entry.param_value + def test_update_sso_settings_keeps_an_env_var_the_read_replica_has_not_caught_up_on( + self, mock_proxy_config, mock_auth, monkeypatch + ): + import json + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( + return_value=SimpleNamespace( + param_value={"SAVED_A_MOMENT_AGO": "new", "MICROSOFT_CLIENT_ID": "old_microsoft_id"} + ) + ) + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=SimpleNamespace(param_value={"MICROSOFT_CLIENT_ID": "old_microsoft_id"}) + ) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + response = client.patch( + "/update/sso_settings", json={"microsoft_client_id": "new_microsoft_id"} + ) + + assert response.status_code == 200 + update_call = mock_prisma.db.litellm_config.update.call_args + assert json.loads(update_call.kwargs["data"]["param_value"]) == {"SAVED_A_MOMENT_AGO": "new"} + def test_get_sso_settings_empty_database( self, mock_proxy_config, mock_auth, monkeypatch ): @@ -2248,7 +2280,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_config.update = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) @@ -2573,7 +2605,7 @@ def test_update_sso_settings_writes_redacted_audit_log(mock_proxy_config, monkey fake_prisma.db.litellm_ssoconfig.upsert = AsyncMock() # No prior SSO row, so before_value resolves to None. fake_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) - fake_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + fake_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) monkeypatch.setattr(proxy_server_module, "premium_user", True) @@ -2645,7 +2677,7 @@ def test_update_sso_settings_audit_captures_redacted_before_snapshot( "google_client_secret": "OLD-SUPER-SECRET", } fake_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=existing_record) - fake_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + fake_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) monkeypatch.setattr(proxy_server_module, "premium_user", True) diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py b/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py index cdaf5c1ee3b..69b54e5c520 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py @@ -72,7 +72,7 @@ def test_config_row_dataclass_shape() -> None: } == { "param_name": "alpha", "param_value": {"k": 1}, - "slots": ("param_name", "param_value"), + "slots": ("param_name", "param_value", "reload_revision"), } @@ -83,12 +83,13 @@ def test_config_row_rejects_unknown_attribute() -> None: def test_pack_config_row_returns_dict_for_caching() -> None: - row = SimpleNamespace(param_name="zeta", param_value=[1, 2, 3]) + row = SimpleNamespace(param_name="zeta", param_value=[1, 2, 3], reload_revision=7) actual = _pack_config_row(row) expanded = {**actual, "is_dict": isinstance(actual, dict)} assert expanded == { "param_name": "zeta", "param_value": [1, 2, 3], + "reload_revision": 7, "is_dict": True, } @@ -119,6 +120,21 @@ def test_unpack_config_row_round_trips_dict() -> None: } +def test_config_row_revision_survives_the_cache_round_trip() -> None: + unpacked = _unpack_config_row( + _pack_config_row(SimpleNamespace(param_name="environment_variables", param_value={}, reload_revision=42)) + ) + assert unpacked is not None and unpacked.reload_revision == 42 + + +@pytest.mark.parametrize("cached_revision", [None, True, "42"]) +def test_unpack_config_row_without_an_integer_revision_has_none(cached_revision: object) -> None: + unpacked = _unpack_config_row( + {"param_name": "environment_variables", "param_value": {}, "reload_revision": cached_revision} + ) + assert unpacked is not None and unpacked.reload_revision is None + + def test_unpack_config_row_error_on_malformed_dict() -> None: with pytest.raises(KeyError): _unpack_config_row({"only_name": "x"}) @@ -168,7 +184,7 @@ async def test_get_config_param_cache_miss_fetches_from_db_and_caches( } assert actual == { "returned": db_row, - "cached": {"param_name": "p2", "param_value": {"y": 2}}, + "cached": {"param_name": "p2", "param_value": {"y": 2}, "reload_revision": None}, "db_called": 1, "db_args": {"key": "param_name", "value": "p2", "table_name": "config"}, } @@ -251,9 +267,9 @@ async def test_prefetch_config_params_populates_cache_for_each_name( "c": _swap_config_cache._store[_config_cache_key("c")], } assert actual == { - "a": {"param_name": "a", "param_value": {"av": 1}}, + "a": {"param_name": "a", "param_value": {"av": 1}, "reload_revision": None}, "b": utils_mod._CONFIG_CACHE_MISS, - "c": {"param_name": "c", "param_value": [3]}, + "c": {"param_name": "c", "param_value": [3], "reload_revision": None}, } diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index e340d4b70ed..163e9611584 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -286,6 +286,151 @@ class TestManagedWebSocketHandlerIntegration: assert captured["model"] == "bedrock_mantle/openai.gpt-5.5" + @pytest.mark.asyncio + async def test_managed_socket_runs_the_frame_guard_on_every_turn(self, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + if "forbidden" in body["input"]: + raise ValueError("blocked by guardrail") + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": body["input"].upper()}), + guardrail_information=({"guardrail_name": "g", "guardrail_cost": 0.5},), + ) + + async def empty_stream(): + return + yield {} + + aresponses: Final = AsyncMock(side_effect=lambda **kwargs: empty_stream()) + monkeypatch.setattr(litellm, "aresponses", aresponses) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-sonnet-4-5", + logging_obj=Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={"user_api_key_hash": "k1"}, + frame_guard=guard, + ) + + await handler._process_response_create(json.dumps({"type": "response.create", "input": "say forbidden"})) + await handler._process_response_create(json.dumps({"type": "response.create", "input": "hello"})) + + aresponses.assert_awaited_once() + error_event: Final = json.loads(websocket.send_text.await_args_list[0].args[0]) + assert error_event["error"] == {"type": "invalid_request_error", "message": "blocked by guardrail"} + sent: Final = aresponses.await_args.kwargs + assert sent["input"] == "HELLO" + assert sent["litellm_metadata"]["user_api_key_hash"] == "k1" + assert sent["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_name": "g", "guardrail_cost": 0.5} + ] + + @pytest.mark.asyncio + async def test_managed_socket_unmasks_the_reply_with_the_frame_guard_s_pii_tokens(self, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler + + presidio = _FakeWSGuardrail() + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + masked = await presidio.check_pii(body["input"], True, None, walk := {}) + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": masked}), pii_tokens=walk["metadata"]["pii_tokens"] + ) + + provider_saw: dict = {} + + async def fake_aresponses(**kwargs): + provider_saw.update(kwargs) + + async def stream(): + yield {"type": "response.output_text.delta", "delta": "Sent to "} + yield { + "type": "response.completed", + "response": { + "id": "resp_1", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "Sent to "}]} + ], + }, + } + + return stream() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-sonnet-4-5", + logging_obj=Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=guard, + guardrail_callbacks=[presidio], + ) + + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "email alice@example.com please"}) + ) + + assert provider_saw["input"] == "email please" + sent: Final = [json.loads(call.args[0]) for call in websocket.send_text.await_args_list] + assert sent[0]["delta"] == "Sent to alice@example.com" + assert sent[1]["response"]["output"][0]["content"][0]["text"] == "Sent to alice@example.com" + + @pytest.mark.asyncio + async def test_http_handler_hands_the_frame_guard_to_a_managed_socket(self): + import json + from unittest.mock import AsyncMock, MagicMock + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + guard: Final = AsyncMock(side_effect=ValueError("blocked by guardrail")) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=ConnectionError("client left")) + + await BaseLLMHTTPHandler().async_responses_websocket( + model="anthropic/claude-sonnet-4-5", + websocket=websocket, + logging_obj=MagicMock(), + responses_api_provider_config=None, + first_message=json.dumps({"type": "response.create", "input": "say forbidden"}), + frame_guard=guard, + ) + + guard.assert_awaited_once() + error_event: Final = json.loads(websocket.send_text.await_args.args[0]) + assert error_event["error"]["message"] == "blocked by guardrail" + @pytest.mark.parametrize("terminal_type", ["response.completed", "response.incomplete"]) @pytest.mark.asyncio async def test_truncated_turn_is_kept_as_history_for_the_next_turn(self, monkeypatch, terminal_type: str): @@ -1321,6 +1466,97 @@ class TestWebSocketProjectQuotaEnforcement: error_event = json.loads(mock_websocket.send_text.call_args[0][0]) assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded") + @pytest.mark.asyncio + async def test_native_handler_runs_the_frame_guard_on_every_turn_and_bills_each_turn_its_own_guardrail_records( + self, + ): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming + + def connection_logging(): + logging_obj = MagicMock() + logging_obj.litellm_params = { + "metadata": { + "standard_logging_guardrail_information": [{"guardrail_name": "admission", "guardrail_cost": 9.0}] + } + } + logging_obj.model_call_details = {} + logging_obj.copy_for_new_call.side_effect = lambda **_: connection_logging() + return logging_obj + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": body["input"] + " [guarded]"}), + guardrail_information=({"guardrail_name": "turn", "guardrail_cost": 0.1 * len(body["input"])},), + ) + + frames = [json.dumps({"type": "response.create", "input": text}) for text in ("hi", "hello")] + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.receive_text = AsyncMock(side_effect=[frames[1], RuntimeError("client closed")]) + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=connection_logging(), + first_message=frames[0], + frame_guard=guard, + ) + await handler.client_to_backend() + + sent = [json.loads(call.args[0]) for call in mock_backend_ws.send.call_args_list] + assert [frame["input"] for frame in sent] == ["hi [guarded]", "hello [guarded]"] + assert [ + turn.logging_obj.litellm_params["metadata"]["standard_logging_guardrail_information"] + for turn in handler._open_turns + ] == [ + [{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.2)}], + [{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.5)}], + ] + + @pytest.mark.asyncio + async def test_native_handler_drops_a_frame_the_guard_rejects_and_keeps_the_socket_open(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming + + async def guard(frame: str) -> GuardedResponsesWsFrame: + if "forbidden" in frame: + raise ValueError("turn guardrail blocked the request") + return GuardedResponsesWsFrame(frame=frame) + + blocked = json.dumps({"type": "response.create", "input": "forbidden"}) + allowed = json.dumps({"type": "response.create", "input": "fine"}) + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.receive_text = AsyncMock(side_effect=[allowed, RuntimeError("client closed")]) + logging_obj = MagicMock() + logging_obj.copy_for_new_call.return_value = logging_obj + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=logging_obj, + first_message=blocked, + frame_guard=guard, + ) + await handler.client_to_backend() + + error_event = json.loads(mock_websocket.send_text.call_args[0][0]) + assert error_event["type"] == "error" + assert error_event["error"] == { + "type": "invalid_request_error", + "message": "turn guardrail blocked the request", + } + assert [json.loads(call.args[0])["input"] for call in mock_backend_ws.send.call_args_list] == ["fine"] + assert len(handler._open_turns) == 1 + @pytest.mark.asyncio async def test_native_handler_forwards_frame_allowed_by_quota_callback(self): from unittest.mock import AsyncMock, MagicMock @@ -2143,6 +2379,34 @@ def _make_streaming(**kwargs): class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" + @pytest.mark.asyncio + async def test_a_reply_to_a_frame_the_guard_masked_is_unmasked_with_the_guard_s_tokens(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame + + guardrail = _FakeWSGuardrail() + + async def guard(frame: str) -> GuardedResponsesWsFrame: + masked = await guardrail.check_pii(json.loads(frame)["input"], True, None, walk := {}) + return GuardedResponsesWsFrame( + frame=json.dumps({"type": "response.create", "input": masked}), + pii_tokens=walk["metadata"]["pii_tokens"], + ) + + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + handler = _make_streaming( + backend_ws=backend_ws, request_data={}, guardrail_callbacks=[guardrail], frame_guard=guard + ) + handler.logging_obj.copy_for_new_call.return_value = handler.logging_obj + + await handler._forward_guarded_frame(json.dumps({"type": "response.create", "input": "email alice@example.com"})) + + assert json.loads(backend_ws.send.await_args.args[0])["input"] == "email " + reply = json.dumps({"type": "response.output_text.delta", "delta": "Sent to "}) + assert json.loads(handler._unmask_response_event(reply))["delta"] == "Sent to alice@example.com" + @pytest.mark.asyncio async def test_mask_response_create_flat_string_input(self): guardrail = _FakeWSGuardrail() diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 001be7e1fbd..04d2b9ddf81 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -19,6 +19,7 @@ Also covers: ErrorEventError.param must accept dict payloads without raising a Pydantic ValidationError (previously typed as Optional[str]). """ +import asyncio import json from datetime import datetime from importlib import import_module @@ -31,15 +32,18 @@ import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, BaseResponsesAPIStreamingIterator, ResponsesAPIStreamingIterator, + ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, _status_code_for_error_fields, ) +from litellm.utils import ProviderConfigManager from litellm.types.llms.openai import ( ErrorEvent, ErrorEventError, @@ -658,3 +662,140 @@ def test_openai_top_level_error_event_keeps_the_provider_code_and_message( assert exc_info.value.status_code == status_code assert type(exc_info.value.original_exception) is exception_type assert f"upstream said {code}" in str(exc_info.value.original_exception) + + +_AZURE_RESPONSES_API_BASE: Final = "https://responses-mid-stream.openai.azure.com" +_AZURE_RESPONSES_REQUEST_URL: Final = f"{_AZURE_RESPONSES_API_BASE}/openai/responses?api-version=2025-04-01-preview" +_AZURE_RESPONSES_FAILED_EVENT: Final = { + "type": "response.failed", + "sequence_number": 2, + "response": { + "id": "resp_mid_stream", + "object": "response", + "created_at": 1790000000, + "status": "failed", + "model": "gpt-5.4-mini", + "output": [], + "error": {"code": "server_error", "message": "upstream failed mid-stream"}, + }, +} +_AZURE_RESPONSES_WS_URL: Final = "wss://responses-mid-stream.openai.azure.com/openai/v1/responses?api-version=preview" + + +def _azure_responses_logging_obj(logged_api_base: str, routed_api_base: str | None) -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.4-mini", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id="responses-mid-stream", + function_id="responses-mid-stream", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-mini", + optional_params={}, + litellm_params=get_litellm_params( + api_base=logged_api_base, + custom_llm_provider="azure", + metadata={ + "model_group": "gpt-5.4-mini-azure", + "deployment": "azure/gpt-5.4-mini", + "api_base": routed_api_base, + }, + ), + ) + return logging_obj + + +def _assert_names_the_azure_deployment(exception: BaseException, api_base: str = _AZURE_RESPONSES_API_BASE) -> None: + debug_info: Final = getattr(exception, "litellm_debug_info", None) + assert isinstance(debug_info, str) + assert f"API Base: `{api_base}`" in debug_info + assert "model_group: `gpt-5.4-mini-azure`" in debug_info + assert "deployment: `azure/gpt-5.4-mini`" in debug_info + assert "api.openai.com" not in debug_info + + +@pytest.mark.asyncio +async def test_response_failed_after_output_reports_the_deployment_it_was_sent_to(): + delta: Final = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_mid_stream", + "output_index": 0, + "content_index": 0, + "delta": "1", + } + body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in (delta, _AZURE_RESPONSES_FAILED_EVENT)) + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body.encode()), + model="gpt-5.4-mini", + responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config( + provider="azure", model="gpt-5.4-mini" + ), + logging_obj=_azure_responses_logging_obj(_AZURE_RESPONSES_REQUEST_URL, _AZURE_RESPONSES_API_BASE), + custom_llm_provider="azure", + ) + stream: Final = aiter(iterator) + assert getattr(await anext(stream), "delta", None) == "1" + + with pytest.raises(MidStreamFallbackError) as exc_info: + await anext(stream) + assert exc_info.value.is_pre_first_chunk is False + _assert_names_the_azure_deployment(exc_info.value.original_exception) + + +class _ScriptedAzureSocket: + def __init__(self, events: tuple[str, ...]) -> None: + self._events = iter(events) + self._drained = asyncio.Event() + + async def recv(self, decode: bool = False) -> str: + event: Final = next(self._events, None) + if event is None: + self._drained.set() + raise ConnectionError("backend closed") + return event + + async def send(self, message: str) -> None: + return None + + async def close(self) -> None: + return None + + async def send_text(self, data: str) -> None: + return None + + async def receive_text(self) -> str: + await self._drained.wait() + raise ConnectionError("client closed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("routed_api_base", "reported_api_base"), + [ + (_AZURE_RESPONSES_API_BASE, _AZURE_RESPONSES_API_BASE), + (None, "wss://responses-mid-stream.openai.azure.com/openai/v1/responses"), + ], + ids=["router", "no_router"], +) +async def test_responses_websocket_turn_failure_reports_the_deployment_it_was_sent_to( + routed_api_base: str | None, reported_api_base: str +): + connection_logging_obj: Final = _azure_responses_logging_obj(_AZURE_RESPONSES_WS_URL, routed_api_base) + connection_logging_obj.pre_call(input=None, api_key="", additional_args={"api_base": _AZURE_RESPONSES_WS_URL}) + socket: Final = _ScriptedAzureSocket((json.dumps(_AZURE_RESPONSES_FAILED_EVENT),)) + streaming: Final = ResponsesWebSocketStreaming( + websocket=socket, + backend_ws=socket, + logging_obj=connection_logging_obj, + authorized_model="gpt-5.4-mini", + custom_llm_provider="azure", + ) + + failure: Final = await streaming.bidirectional_forward() + + assert failure is not None + _assert_names_the_azure_deployment(failure, reported_api_base) diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 018fa808f70..e5565a3df16 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -15,10 +15,12 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, + _honors_one_hour_cache_ttl, _get_min_token_count_for_deployments, ) from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter MODEL_GROUP_ALIAS = "my-claude-group" @@ -126,6 +128,21 @@ def test_get_min_token_count_for_deployments_falls_back_to_default_for_empty_gro assert _get_min_token_count_for_deployments([]) == DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT +def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> None: + standard_logging_object: Final = cast( + StandardLoggingPayload, + { + "custom_llm_provider": "bedrock", + "model": "bedrock/converse/arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile", + }, + ) + + assert _honors_one_hour_cache_ttl( + standard_logging_object, + {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + ) is True + + @pytest.mark.asyncio async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minimum(): """ diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 56184fa299a..637fdfa6873 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -14,6 +14,7 @@ import respx import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import FakeAnthropicMessagesStreamIterator from litellm.litellm_core_utils import get_llm_provider_logic from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import ( @@ -32,6 +33,7 @@ from litellm.router_utils.fallback_event_handlers import ( is_bad_request_without_fallback, log_failure_fallback_event, log_success_fallback_event, + mid_stream_fallback_hop_kwargs, mid_stream_fallback_snapshot_kwargs, mid_stream_retry_kwargs, record_pre_routing_selection, @@ -1865,6 +1867,41 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through(): assert len(fake.calls) == 1 +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.asyncio +async def test_refused_stream_built_from_a_finished_agentic_turn_is_delivered_without_a_fallback(): + """Websearch, MCP and advisor loops run their tools before streaming the final reply, so a + fallback would run every tool call again for a reply the client then gets from another model.""" + fake: Final = FakeAnthropicUpstream() + router: Final = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream: Final = await router._aanthropic_messages_streaming_iterator( + response=FakeAnthropicMessagesStreamIterator(REFUSAL_RESPONSE), + initial_kwargs=mid_stream_fallback_hop_kwargs( + model="fable-tier", + original_generic_function=litellm.anthropic_messages, + controls=None, + kwargs={ + "model": "fable-tier", + "max_tokens": 16, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ), + ) + events: Final = [ + json.loads(line[len(b"data: ") :]) + for line in (await _collect(stream)).splitlines() + if line.startswith(b"data: ") + ] + + assert fake.calls == [] + message_delta: Final = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"] == REFUSAL_RESPONSE["stop_details"] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index c61ebb3737c..2d13829bf33 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -2,7 +2,7 @@ import json import asyncio import datetime import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from types import MappingProxyType, SimpleNamespace from typing import Final, Literal, cast @@ -17,6 +17,10 @@ from litellm import TranscriptionResponse, model_cost from litellm.constants import REALTIME_METERED_USAGE_KEY from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.fallback_generalizations import ( + get_fallback_generalization_rules, + set_fallback_generalizations, +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -75,6 +79,48 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +@pytest.fixture +def _claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-baseline", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": {"mode": "chat", "max_input_tokens": 200000, "max_output_tokens": 64000}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +@pytest.fixture +def _priced_claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-priced", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": { + "mode": "chat", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "input_cost_per_token": 0.0000321, + "output_cost_per_token": 0.0000765, + }, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch): """ Router/proxy configs may use deployment ids like openai/openai/. Cost lookup must @@ -149,6 +195,265 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_completion_cost_prefers_exact_route_price_over_capability_generalization( + _local_model_cost_map: None, + _claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + priced_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + monkeypatch.setattr(litellm, "model_cost", {priced_model: priced_entry}) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=usage, + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + ) + + input_rate: Final = priced_entry["input_cost_per_token"] + output_rate: Final = priced_entry["output_cost_per_token"] + assert isinstance(input_rate, (int, float)), priced_entry + assert isinstance(output_rate, (int, float)), priced_entry + assert cost == usage.prompt_tokens * input_rate + usage.completion_tokens * output_rate + + +def test_completion_cost_keeps_priced_generalization_before_later_exact_candidate( + _local_model_cost_map: None, + _priced_claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + exact_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + monkeypatch.setattr(litellm, "model_cost", {priced_model: exact_entry}) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=usage, + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000321 + usage.completion_tokens * 0.0000765) + + +@pytest.fixture +def _per_second_priced_claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-per-second", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": {"cost_per_second": 0.25}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_completion_cost_keeps_per_second_priced_generalization_before_later_exact_candidate( + _local_model_cost_map: None, + _per_second_priced_claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + monkeypatch.setattr( + litellm, + "model_cost", + { + priced_model: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + }, + ) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + total_time=2000, + ) + + assert cost == pytest.approx(0.25 * 2) + + +@pytest.fixture +def _price_free_gpt_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "gpt-family-baseline", + "pattern": r"gpt-\d+(?:\.\d+)?", + "model_info": {"mode": "chat", "supports_reasoning": True}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_completion_cost_prices_rule_only_response_model_at_the_azure_deployment_rate( + _local_model_cost_map: None, + _price_free_gpt_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + deployment_model: Final = "azure/gpt-4o-2024-11-20" + monkeypatch.setattr( + litellm, + "model_cost", + { + deployment_model: { + "litellm_provider": "azure", + "mode": "chat", + "input_cost_per_token": 0.0000111, + "output_cost_per_token": 0.0000222, + } + }, + ) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse(model="gpt-9.9", choices=[], usage=usage) + + cost: Final = completion_cost( + completion_response=response, + model=deployment_model, + custom_llm_provider="azure", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222) + + +def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( + _local_model_cost_map: None, +) -> None: + model: Final = "lazy-candidate-regression-org/lazy-candidate-model" + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model=model, + choices=[], + usage=usage, + ) + + with respx.mock(assert_all_called=False) as huggingface: + huggingface.get(host="huggingface.co").mock( + return_value=httpx.Response(200, json={"max_position_embeddings": 4096}) + ) + cost: Final = completion_cost( + completion_response=response, + model=f"huggingface/{model}", + custom_llm_provider="huggingface", + ) + + assert len(huggingface.calls) == 1 + assert cost == 0.0 + + +def test_completion_cost_skips_huggingface_config_request_with_custom_pricing( + _local_model_cost_map: None, +) -> None: + model: Final = "unmapped-cost-regression-org/unmapped-model" + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model=model, + choices=[], + usage=usage, + ) + + with respx.mock(assert_all_called=False) as huggingface: + huggingface.get(host="huggingface.co").mock( + return_value=httpx.Response(200, json={"max_position_embeddings": 4096}) + ) + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="huggingface", + custom_pricing=True, + custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + + assert len(huggingface.calls) == 0 + assert cost == pytest.approx(usage.prompt_tokens * 0.001 + usage.completion_tokens * 0.002) + + +def test_completion_cost_returns_zero_for_capability_only_generalization( + _local_model_cost_map: None, + _claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model: Final = "bedrock/invoke/anthropic.claude-haiku-9-9-20991231-v1:0" + monkeypatch.setattr(litellm, "model_cost", {}) + response: Final = ModelResponse( + model="anthropic.claude-haiku-9-9-20991231-v1:0", + choices=[], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + + assert cost == 0.0 + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "hosted_vllm"]) +def test_completion_cost_prices_base_model_under_its_own_provider( + custom_llm_provider: str, _local_model_cost_map: None +) -> None: + response: Final = ModelResponse( + model="my-deploy", + choices=[], + usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000), + ) + + cost: Final = completion_cost( + completion_response=response, + model=f"{custom_llm_provider}/my-deploy", + custom_llm_provider=custom_llm_provider, + base_model="claude-sonnet-4-5", + ) + + info: Final = litellm.get_model_info("claude-sonnet-4-5", custom_llm_provider="anthropic") + assert cost == pytest.approx( + 1000 * info["input_cost_per_token"] + 1000 * info["output_cost_per_token"] + ) + + def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: dated_response = ModelResponse( model="gpt-5.6-luna-2099-01-01", @@ -229,6 +534,62 @@ def test_response_cost_calculator_custom_pricing_ignores_upstream_response_cost_ assert cost == pytest.approx(expected_cost) +def _header_cost_chat_response(additional_headers: dict[str, str]) -> ModelResponse: + response: Final = ModelResponse( + model="cc-header-cost-chat", + choices=[Choices(message=Message(content="ok"))], + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + response._hidden_params = {"additional_headers": additional_headers} + return response + + +def _cost_with_upstream_headers(monkeypatch: pytest.MonkeyPatch, additional_headers: dict[str, str]) -> float: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.register_model( + model_cost={ + "cc-header-cost-chat": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + return response_cost_calculator( + response_object=_header_cost_chat_response(additional_headers), + model="cc-header-cost-chat", + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + +@pytest.mark.parametrize("header_cost", ("-5", "not-a-number", "nan", "inf")) +def test_response_cost_calculator_prices_from_the_map_when_upstream_cost_header_is_not_billable( + monkeypatch: pytest.MonkeyPatch, + header_cost: str, +) -> None: + price_map_cost: Final = _cost_with_upstream_headers(monkeypatch, {}) + + cost: Final = _cost_with_upstream_headers(monkeypatch, {"llm_provider-x-litellm-response-cost": header_cost}) + + assert price_map_cost > 0 + assert cost == pytest.approx(price_map_cost) + + +@pytest.mark.parametrize(("header_cost", "expected_cost"), (("0", 0.0), ("0.25", 0.25))) +def test_response_cost_calculator_bills_a_finite_non_negative_upstream_cost_header( + monkeypatch: pytest.MonkeyPatch, + header_cost: str, + expected_cost: float, +) -> None: + cost: Final = _cost_with_upstream_headers(monkeypatch, {"llm_provider-x-litellm-response-cost": header_cost}) + + assert cost == expected_cost + + def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): class MockResponse(BaseModel): pass diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index 1b9c6950ba9..a705597952b 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -22,7 +22,6 @@ def _load(name: str) -> ModuleType: return module -schema_module: Final = _load("generate_model_prices_schema") guard: Final = _load("cost_map_guard") MAP_FILES: Final = (guard.COST_MAP_PATH,) @@ -52,12 +51,9 @@ def _serialize(cost_map: dict[str, object]) -> str: return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" -def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object: +def _snapshot(cost_map: dict[str, object], backup: str | None = None) -> object: text = _serialize(cost_map) - rendered = schema_module.render(schema_module.build_schema(cost_map)) - return guard.Snapshot( - cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema - ) + return guard.Snapshot(cost_map=text, backup=text if backup is None else backup) BASE: Final = _snapshot(BASE_MAP) @@ -133,14 +129,14 @@ def test_bot_may_add_and_reprice_models() -> None: def test_broken_json_is_reported() -> None: - head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}") + head = guard.Snapshot(cost_map="{not json", backup="{not json") assert _failures(head, bot=False) == ( f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", ) def test_non_object_root_is_reported() -> None: - head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}") + head = guard.Snapshot(cost_map="[]", backup="[]") assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",) @@ -149,28 +145,7 @@ def test_backup_drift_is_reported() -> None: assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)] -def test_schema_out_of_sync_is_reported() -> None: - head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema) - assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)] - - -def test_schema_validation_errors_are_reported() -> None: - head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)}) - prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c." - assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [ - "input_cost_per_token", - "output_cost_per_token", - ] - - -def test_unclassified_entry_key_is_reported() -> None: - text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)}) - head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema) - (failure,) = _failures(head, bot=False) - assert "Unclassified keys" in failure and "weird_thing" in failure - - -STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}), schema="{}") +STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) CODE_ONLY: Final = ( "litellm/utils.py", "tests/test_litellm/test_utils.py", @@ -179,7 +154,7 @@ CODE_ONLY: Final = ( def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None: - unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", schema="") + unparseable: Final = guard.Snapshot(cost_map="{not json", backup="") assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == () assert _failures(STALE_HEAD, changed_files=(), bot=False) == () assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == () @@ -189,18 +164,16 @@ def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None: failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] - assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] def test_bot_pr_always_gets_the_file_checks() -> None: failures: Final = _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=True) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] - assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] def test_contract_names_the_skip() -> None: assert guard.contract_for(False, CODE_ONLY) == "human PR, cost map untouched" - assert guard.contract_for(False, (*CODE_ONLY, guard.SCHEMA_PATH)) == "human PR, file checks only" + assert guard.contract_for(False, (*CODE_ONLY, guard.COST_MAP_PATH)) == "human PR, file checks only" assert guard.contract_for(True, CODE_ONLY) == "bot contract enforced" @@ -213,6 +186,13 @@ def test_bot_may_only_touch_the_cost_map_files() -> None: ) +def test_bot_may_update_cost_map_and_legacy_schema() -> None: + head: Final = _snapshot({**BASE_MAP, "openrouter/c": _entry()}) + changed: Final = (*guard.GUARDED_PATHS, guard.LEGACY_SCHEMA_PATH) + + assert _failures(head, changed_files=changed) == () + + def test_bot_may_not_remove_models() -> None: head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}) assert _failures(head, bot=False) == () @@ -236,7 +216,6 @@ def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: (repo / guard.COST_MAP_PATH).write_text(text) (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) (repo / guard.BACKUP_PATH).write_text(text) - (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map))) return _git_commit(repo, message) @@ -290,8 +269,7 @@ def test_main_skips_the_file_checks_on_a_stale_base_the_pr_never_touched(tmp_pat subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) _commit(tmp_path, BASE_MAP, "base") (tmp_path / guard.BACKUP_PATH).write_text(_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) - (tmp_path / guard.SCHEMA_PATH).write_text("{}") - stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup and schema") + stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup") head: Final = _commit_code_only(tmp_path, "code change on the stale base") human: Final = _run_guard(tmp_path, stale_base, head, "litellm_fix_pricing") assert human.returncode == 0, human.stdout + human.stderr diff --git a/tests/unit/test_litellm_cli.py b/tests/unit/test_litellm_cli.py new file mode 100644 index 00000000000..5d1e510742a --- /dev/null +++ b/tests/unit/test_litellm_cli.py @@ -0,0 +1,344 @@ +import json +import os +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from litellm_cli.main import cli + +_HELP_ARGV: Final = ( + ("--help",), + ("logout", "--help"), + ("login", "--help"), + ("models", "list", "--help"), + ("claude", "--help"), + ("opencode", "--help"), + ("codex", "--help"), + ("pi", "--help"), + ("keys", "list", "--help"), +) + +_COMMAND_ARGV: Final = (("logout",), ("whoami",)) + +_RUN_EACH_ARGV: Final = """ +import contextlib, io, json, sys +from {module} import cli +runs = {{}} +for argv in json.loads(sys.argv[1]): + sys.argv = ["lite", *argv] + output = io.StringIO() + with contextlib.redirect_stdout(output): + try: + cli() + except SystemExit as stopped: + code = stopped.code + else: + code = 0 + runs[" ".join(argv)] = {{"exit": code, "output": output.getvalue()}} +print(json.dumps({{"runs": runs, "sdk_initialized": "completion" in vars(sys.modules["litellm"])}})) +""" + +_USE_SDK_AFTER_HELP: Final = """ +import builtins, importlib, json, sys +interpreter_imports = (builtins.__import__, importlib.import_module) +from {module} import cli +sys.argv = ["lite", "--help"] +try: + cli() +except SystemExit: + pass +before = "completion" in vars(sys.modules["litellm"]) +{first_use} +from litellm.router import Router +import litellm +import litellm.responses.main +from litellm import acompletion +print(json.dumps({{ + "before": before, + "router": Router.__name__, + "completion": callable(litellm.completion), + "acompletion": acompletion is litellm.acompletion, + "responses_is_function": callable(litellm.responses) and not isinstance(litellm.responses, type(litellm)), + "namespace": sorted(vars(litellm)), + "module_type": type(litellm).__name__, + "imports_restored": (builtins.__import__, importlib.import_module) == interpreter_imports, +}})) +""" + +_FIRST_USE_FROM_THREADS: Final = """ +import json, sys, threading +from litellm_cli.main import cli +sys.argv = ["lite", "--help"] +try: + cli() +except SystemExit: + pass +import litellm +start = threading.Barrier(4) +seen = [] +def use_sdk(): + start.wait() + try: + seen.append(callable(litellm.completion)) + except AttributeError as error: + seen.append(str(error)) +threads = [threading.Thread(target=use_sdk) for _ in range(4)] +for thread in threads: + thread.start() +for thread in threads: + thread.join() +print(json.dumps({"seen": seen})) +""" + +_STAR_IMPORT_AFTER_HELP: Final = """ +import json, sys +from litellm_cli.main import cli +sys.argv = ["lite", "--help"] +try: + cli() +except SystemExit: + pass +before = "completion" in vars(sys.modules["litellm"]) +from litellm import * +print(json.dumps({"before": before, "completion": callable(globals().get("completion")), "acompletion": callable(globals().get("acompletion"))})) +""" + +_SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING: Final = { + "sdk_sync.py": """ +import threading +initializing = threading.Event() +submodule_loading = threading.Event() +""", + "litellm/__init__.py": """ +import sdk_sync +sdk_sync.initializing.set() +sdk_sync.submodule_loading.wait(2) +from litellm import submodule +ready = True +""", + "litellm/submodule.py": """ +import litellm +import sdk_sync +sdk_sync.submodule_loading.set() +ready_when_loaded = getattr(litellm, "ready", False) +""", + "litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n", + "litellm/proxy/__init__.py": "", + "litellm/proxy/client/__init__.py": "", + "litellm/proxy/client/cli/__init__.py": "", + "litellm/proxy/client/cli/main.py": """ +import json, threading +import litellm +import sdk_sync + +def cli(): + def initialize_sdk(): + litellm.ready + + def import_submodule_once_initialization_starts(): + sdk_sync.initializing.wait() + import litellm.submodule + + threads = [threading.Thread(target=initialize_sdk, daemon=True), + threading.Thread(target=import_submodule_once_initialization_starts, daemon=True)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(10) + print(json.dumps({"finished": [not thread.is_alive() for thread in threads]})) + +litellm_proxy_cli = cli +""", +} + +_SDK_WHOSE_INIT_FAILS: Final = { + "litellm/__init__.py": 'raise ImportError("a dependency of the SDK is broken")\n', + "litellm/utils.py": "", + "litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n", + "litellm/proxy/__init__.py": "", + "litellm/proxy/client/__init__.py": "", + "litellm/proxy/client/cli/__init__.py": "", + "litellm/proxy/client/cli/main.py": """ +import json +import litellm + +def error_from(use): + try: + use() + except Exception as error: + return f"{type(error).__name__}: {error}" + return None + +def import_sdk_code(): + import litellm.utils + +def cli(): + print(json.dumps([error_from(lambda: litellm.completion), error_from(lambda: litellm.completion), error_from(import_sdk_code)])) + +litellm_proxy_cli = cli +""", +} + +_RUN_LITE_AGAINST_SDK_AT: Final = """ +import sys +sys.path.insert(0, sys.argv[1]) +from litellm_cli.main import cli +cli() +""" + +_COMMAND_MODULES_LOADED_BY: Final = """ +import json, sys +from litellm_cli.main import cli +sys.argv = ["lite", *sys.argv[1:]] +try: + cli() +except SystemExit: + pass +prefix = "litellm.proxy.client.cli.commands." +print(json.dumps({"modules": sorted(name[len(prefix):] for name in sys.modules if name.startswith(prefix))})) +""" + + +def _run(source: str, argv: tuple[str, ...], home: Path) -> dict[str, object]: + result: Final = subprocess.run( + [sys.executable, "-I", "-c", source, *argv], + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "HOME": str(home), "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + assert result.returncode == 0, result.stderr + return json.loads(result.stdout.rstrip("\n").rpartition("\n")[2]) + + +@pytest.fixture(scope="module") +def cli_runs(tmp_path_factory: pytest.TempPathFactory) -> Mapping[str, dict[str, object]]: + argv: Final = (json.dumps(_HELP_ARGV + _COMMAND_ARGV),) + return { + "thin": _run(_RUN_EACH_ARGV.format(module="litellm_cli.main"), argv, tmp_path_factory.mktemp("thin")), + "full": _run( + _RUN_EACH_ARGV.format(module="litellm.proxy.client.cli.main"), argv, tmp_path_factory.mktemp("full") + ), + } + + +def test_lite_runs_commands_without_initializing_the_sdk(cli_runs: Mapping[str, dict[str, object]]): + assert cli_runs["thin"]["sdk_initialized"] is False + assert cli_runs["full"]["sdk_initialized"] is True + + +@pytest.mark.parametrize("argv", _HELP_ARGV + _COMMAND_ARGV, ids=" ".join) +def test_lite_prints_what_the_full_sdk_cli_prints(argv: tuple[str, ...], cli_runs: Mapping[str, dict[str, object]]): + thin: Final = cli_runs["thin"]["runs"] + full: Final = cli_runs["full"]["runs"] + assert isinstance(thin, dict) and isinstance(full, dict) + + assert thin[" ".join(argv)] == full[" ".join(argv)] + assert thin[" ".join(argv)]["output"] + + +@pytest.mark.parametrize( + ("argv", "loaded", "not_loaded"), + [ + (("logout", "--help"), {"auth"}, {"agents", "autoroute.commands", "models", "keys"}), + (("models", "list", "--help"), {"models"}, {"agents", "autoroute.commands", "keys"}), + (("claude", "--help"), {"agents"}, {"autoroute.commands", "models", "keys"}), + ], +) +def test_lite_imports_only_the_command_it_runs( + argv: tuple[str, ...], loaded: set[str], not_loaded: set[str], tmp_path: Path +): + modules: Final = _run(_COMMAND_MODULES_LOADED_BY, argv, tmp_path)["modules"] + assert isinstance(modules, list) + + assert loaded <= set(modules) + assert not not_loaded & set(modules) + + +@pytest.fixture(scope="module") +def sdk_used_after_full_cli(tmp_path_factory: pytest.TempPathFactory) -> dict[str, object]: + return _run( + _USE_SDK_AFTER_HELP.format(module="litellm.proxy.client.cli.main", first_use="pass"), + (), + tmp_path_factory.mktemp("full"), + ) + + +@pytest.mark.parametrize( + "first_use", + [ + "import litellm; litellm.completion", + "import litellm; litellm.types.utils.ModelResponse", + "from litellm import acompletion", + "from litellm.router import Router", + "from litellm.types import utils", + "import litellm.responses.main", + 'importlib.import_module("litellm.router")', + 'importlib.import_module(".utils", "litellm")', + 'exec("from .types import utils", {"__name__": "litellm.lazy_import", "__package__": "litellm"})', + ], +) +def test_sdk_initializes_in_place_on_first_use_after_lite_imports( + first_use: str, sdk_used_after_full_cli: dict[str, object], tmp_path: Path +): + thin: Final = _run(_USE_SDK_AFTER_HELP.format(module="litellm_cli.main", first_use=first_use), (), tmp_path) + + assert thin["before"] is False + assert sdk_used_after_full_cli["before"] is True + assert {**thin, "before": True} == sdk_used_after_full_cli + assert (thin["router"], thin["completion"], thin["acompletion"], thin["responses_is_function"]) == ( + "Router", + True, + True, + True, + ) + + +def test_star_import_after_lite_imports_gets_the_whole_sdk(tmp_path: Path): + assert _run(_STAR_IMPORT_AFTER_HELP, (), tmp_path) == {"before": False, "completion": True, "acompletion": True} + + +def test_threads_that_use_the_sdk_at_once_all_wait_for_one_initialization(tmp_path: Path): + used: Final = _run(_FIRST_USE_FROM_THREADS, (), tmp_path) + + assert used == {"seen": [True, True, True, True]} + + +def _write_sdk(root: Path, files: Mapping[str, str]) -> Path: + for relative_path, source in files.items(): + (root / relative_path).parent.mkdir(parents=True, exist_ok=True) + (root / relative_path).write_text(source) + return root + + +def test_importing_sdk_code_while_another_thread_initializes_the_sdk_does_not_hang(tmp_path: Path): + sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING) + + assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == {"finished": [True, True]} + + +def test_every_use_after_a_failed_sdk_init_raises_the_init_error(tmp_path: Path): + sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_INIT_FAILS) + + assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == ["ImportError: a dependency of the SDK is broken"] * 3 + + +def test_lite_keeps_an_sdk_that_was_already_imported( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +): + monkeypatch.setitem(sys.modules, "litellm", litellm) + monkeypatch.setattr(litellm, "success_callback", ["kept"]) + monkeypatch.setattr(sys, "argv", ["lite", "--help"]) + + with pytest.raises(SystemExit) as exited: + cli() + + assert exited.value.code == 0 + assert "LiteLLM Proxy CLI - Manage your LiteLLM proxy server" in capsys.readouterr().out + assert sys.modules["litellm"] is litellm + assert litellm.success_callback == ["kept"] diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 14cea8f0c2c..0293af05c84 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -21,15 +21,22 @@ import httpx import pytest import respx from openai import APITimeoutError, AsyncOpenAI +from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion import litellm -from litellm import acompletion, completion -from litellm import acompletion_with_retries, aresponses_with_retries, completion_with_retries, responses_with_retries +from litellm import ( + acompletion, + acompletion_with_retries, + aresponses_with_retries, + completion, + completion_with_retries, + responses_with_retries, +) from litellm import main as litellm_main -from litellm.constants import CONTROL_OPTIONS_KEY from litellm.caching.base_cache import BaseCache from litellm.caching.caching import Cache +from litellm.constants import CONTROL_OPTIONS_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs @@ -130,7 +137,7 @@ def test_completion_missing_role(openai_api_response): client = OpenAI(api_key="test_api_key") - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = { "x-request-id": "123", "openai-organization": "org-123", @@ -142,7 +149,7 @@ def test_completion_missing_role(openai_api_response): print(f"openai_api_response: {openai_api_response}") with patch.object( - client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) + client, "post", MagicMock(return_value=mock_raw_response) ) as mock_create: litellm.completion( model="gpt-4o-mini", @@ -218,7 +225,7 @@ async def test_url_with_format_param_openai(model, sync_mode): } ], } - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: if sync_mode: response = completion(**args, client=client) @@ -232,7 +239,7 @@ async def test_url_with_format_param_openai(model, sync_mode): print(mock_client.call_args.kwargs) - json_str = json.dumps(mock_client.call_args.kwargs) + json_str = json.dumps(mock_client.call_args.kwargs["body"]) assert "format" not in json_str @@ -3005,9 +3012,10 @@ def test_stream_chunk_builder_thinking_blocks(): assert response.choices[0].message.thinking_blocks is not None -from litellm.llms.openai.openai import OpenAIChatCompletion import traceback +from litellm.llms.openai.openai import OpenAIChatCompletion + user_message = "Write a short poem about the sky" messages = [{"content": user_message, "role": "user"}] @@ -3777,7 +3785,7 @@ def test_completion_forwards_store_and_prompt_cache_key_to_openai(): client = OpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: litellm.completion( model="openai/gpt-4o", @@ -3790,7 +3798,7 @@ def test_completion_forwards_store_and_prompt_cache_key_to_openai(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["store"] is False assert request_body["prompt_cache_key"] == "test-cache-key" @@ -3805,7 +3813,7 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): client = AsyncOpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: await litellm.acompletion( model="openai/gpt-4o", @@ -3818,7 +3826,7 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["store"] is False assert request_body["prompt_cache_key"] == "test-cache-key" @@ -3832,7 +3840,7 @@ def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): client = OpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: litellm.completion( model="openai/gpt-4o", @@ -3843,7 +3851,7 @@ def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert "store" not in request_body assert "prompt_cache_key" not in request_body @@ -4038,8 +4046,7 @@ def test_openai_model_without_a_provider_still_routes_to_openai(): from openai import OpenAI client = OpenAI(api_key="fake-key") - raw_response = client.chat.completions.with_raw_response - with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): + with patch.object(client, "post") as mock_create, contextlib.suppress(Exception): litellm.completion( model="gpt-4o", messages=[{"role": "user", "content": "hello"}], @@ -4050,7 +4057,7 @@ def test_openai_model_without_a_provider_still_routes_to_openai(): def _openai_chat_create_kwargs(client, **completion_kwargs): - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: with contextlib.suppress(Exception): litellm.completion( messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], @@ -4075,7 +4082,7 @@ def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6 from openai import OpenAI client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1")["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4087,7 +4094,7 @@ def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6 from openai import OpenAI client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1")["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4100,7 +4107,7 @@ async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_ from openai import AsyncOpenAI client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: with contextlib.suppress(Exception): await litellm.acompletion( model="gpt-5.6", @@ -4111,7 +4118,7 @@ async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_ ) mock_create.assert_called_once() - request_body = mock_create.call_args.kwargs + request_body = mock_create.call_args.kwargs["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4123,12 +4130,12 @@ def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6() from openai import OpenAI client = OpenAI(api_key="fake-api-key") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") + request = _openai_chat_create_kwargs(client, model="gpt-5.6") - assert request_body["messages"][0]["content"] == [ + assert request["body"]["messages"][0]["content"] == [ {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} ] - assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "implicit"} + assert request["options"]["extra_json"]["prompt_cache_options"] == {"mode": "implicit"} _SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" @@ -4556,7 +4563,7 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging logging_obj: Final = LiteLLMLogging( - model="us.anthropic.claude-opus-5", + model="us.anthropic.claude-unpriced-999", messages=[{"role": "user", "content": "hi"}], stream=True, call_type="completion", @@ -4582,6 +4589,41 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): assert response._hidden_params.get("response_cost") is None +def test_stream_chunk_builder_prices_alias_response_at_its_deployment_model(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + deployment_model: Final = "us.anthropic.claude-opus-5" + logging_obj: Final = LiteLLMLogging( + model=deployment_model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-priced", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + entry: Final = litellm.model_cost[deployment_model] + expected_cost: Final = 40 * entry["input_cost_per_token"] + 5 * entry["output_cost_per_token"] + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(expected_cost) + assert response._hidden_params.get("response_cost") == pytest.approx(expected_cost) + + def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) @@ -5141,6 +5183,7 @@ def test_drop_params_false_still_rejects_an_invalid_stream_chunk_size() -> None: def test_acompletion_params(): import inspect + from litellm.types.completion import CompletionRequest acompletion_params_odict = inspect.signature(acompletion).parameters @@ -5171,7 +5214,7 @@ def test_acompletion_params(): def _openai_mock_response(*args: object, **kwargs: object) -> MagicMock: - new_response: Final = MagicMock() + new_response: Final = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} response_object: Final = { "id": "chatcmpl-123", @@ -5206,7 +5249,7 @@ def test_null_role_response(): openai_client = openai.OpenAI() with patch.object( - openai_client.chat.completions, "create", side_effect=_openai_mock_response + openai_client, "post", side_effect=_openai_mock_response ) as mock_response: response = litellm.completion( model="gpt-3.5-turbo", @@ -5296,7 +5339,7 @@ def test_completion_perplexity_api(): pydantic_obj = ChatCompletion(**response_object) def _return_pydantic_obj(*args, **kwargs): - new_response = MagicMock() + new_response = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} new_response.parse.return_value = pydantic_obj @@ -5305,8 +5348,7 @@ def test_completion_perplexity_api(): openai_client = OpenAI() with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", side_effect=_return_pydantic_obj, ) as mock_client: # litellm.set_verbose= True @@ -5350,7 +5392,7 @@ async def test_openai_compatible_custom_api_base(provider): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -5368,7 +5410,7 @@ async def test_openai_compatible_custom_api_base(provider): print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) - assert "hello" in mock_call.call_args.kwargs["extra_body"] + assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"] @pytest.mark.parametrize( @@ -5401,7 +5443,7 @@ async def test_openai_compatible_custom_api_video(provider): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -7264,7 +7306,7 @@ def _openai_hallucinated_tool_call_mock_response( *args: object, **kwargs: object, ) -> MagicMock: - new_response: Final = MagicMock() + new_response: Final = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} response_object: Final = { "id": "chatcmpl-123", @@ -7322,8 +7364,8 @@ def test_openai_hallucinated_tool_call(): openai_client = openai.OpenAI() with patch.object( - openai_client.chat.completions, - "create", + openai_client, + "post", side_effect=_openai_hallucinated_tool_call_mock_response, ) as mock_response: response = litellm.completion( @@ -7445,7 +7487,7 @@ def test_completion_novita_ai(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions.with_raw_response, "create" + openai_client, "post", return_value=MagicMock(spec=LegacyAPIResponse) ) as mock_call: mock_call.return_value.headers = {} mock_call.return_value.parse.return_value = litellm.ModelResponse( @@ -7464,11 +7506,11 @@ def test_completion_novita_ai(): # Verify model is passed correctly assert ( - mock_call.call_args.kwargs["model"] + mock_call.call_args.kwargs["body"]["model"] == "meta-llama/llama-3.3-70b-instruct" ) # Verify messages are passed correctly - assert mock_call.call_args.kwargs["messages"] == messages + assert mock_call.call_args.kwargs["body"]["messages"] == messages except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -7491,8 +7533,8 @@ def test_completion_novita_ai_dynamic_params(api_key): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, - "create", + openai_client, + "post", side_effect=Exception("Invalid API key"), ) as mock_call: with pytest.raises(Exception, match="Invalid API key") as exc_info: @@ -7524,7 +7566,7 @@ def test_completion_openai_metadata(monkeypatch, enable_preview_features): monkeypatch.setattr(litellm, "enable_preview_features", enable_preview_features) with patch.object( - client.chat.completions.with_raw_response, "create", return_value=MagicMock() + client, "post", return_value=MagicMock() ) as mock_completion: try: resp = litellm.completion( @@ -7538,11 +7580,11 @@ def test_completion_openai_metadata(monkeypatch, enable_preview_features): mock_completion.assert_called_once() if enable_preview_features: - assert mock_completion.call_args.kwargs["metadata"] == { + assert mock_completion.call_args.kwargs["body"]["metadata"] == { "my-test-key": "my-test-value" } else: - assert "metadata" not in mock_completion.call_args.kwargs + assert "metadata" not in mock_completion.call_args.kwargs["body"] AZURE_TTS_BASE: Final = "https://tts.example.azure.com" diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_and_context_window.py similarity index 81% rename from tests/unit/test_model_prices_schema.py rename to tests/unit/test_model_prices_and_context_window.py index ea484c4a2f4..395fe8c9750 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -1,6 +1,5 @@ from __future__ import annotations -import importlib.util import json import re from collections.abc import Mapping @@ -8,7 +7,6 @@ from pathlib import Path from types import MappingProxyType from typing import Final -import jsonschema import pytest from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter @@ -19,27 +17,8 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts REPO_ROOT = Path(__file__).parents[2] -GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py" PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" BACKUP_PRICES_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" - - -def build_validator(schema: dict) -> jsonschema.Draft202012Validator: - return jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER) - - -def load_generator(): - spec = importlib.util.spec_from_file_location("generate_model_prices_schema", GENERATOR_PATH) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -@pytest.fixture(scope="module") -def committed_schema() -> dict: - return json.loads(SCHEMA_PATH.read_text()) @pytest.fixture(scope="module") @@ -168,170 +147,6 @@ def test_github_copilot_messages_claude_rows_keep_their_anthropic_thinking_capab } -def test_committed_schema_matches_generator_output(prices: dict, committed_schema: dict): - generator = load_generator() - regenerated = json.loads(generator.render(generator.build_schema(prices))) - assert regenerated == committed_schema, ( - "model_prices_and_context_window.schema.json is out of sync; " - "run `python ci_cd/generate_model_prices_schema.py` and commit the result" - ) - - -def test_prices_file_validates_against_committed_schema(prices: dict, committed_schema: dict): - validator = build_validator(committed_schema) - errors = [ - f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}" - for error in validator.iter_errors(prices) - ] - assert errors == [] - - -@pytest.mark.parametrize( - "entry", - [ - {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": "0.01"}, - {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": -1}, - {"litellm_provider": "openai", "mode": "not_a_real_mode"}, - {"mode": "chat"}, - {"litellm_provider": "openai", "deprecation_date": "June 2026"}, - {"litellm_provider": "openai", "deprecation_date": "2026-99-99"}, - {"litellm_provider": "openai", "deprecation_date": "2026-13-01"}, - {"litellm_provider": "openai", "deprecation_date": "2026-01-32"}, - {"litellm_provider": "openai", "deprecation_date": "2026-01-00"}, - {"litellm_provider": "openai", "deprecation_date": "2026-02-31"}, - {"litellm_provider": "openai", "supported_modalities": ["smell"]}, - {"litellm_provider": "openai", "supports_vision": "yes"}, - {"litellm_provider": "openai", "max_tokens": 8191.5}, - {"litellm_provider": "openai", "tiered_pricing": [{"unknown_tier_field": 1}]}, - ], - ids=[ - "cost_as_string", - "negative_cost", - "unknown_mode", - "missing_provider", - "non_iso_deprecation_date", - "impossible_month_and_day", - "month_out_of_range", - "day_out_of_range", - "day_zero", - "calendar_impossible_day", - "unknown_modality", - "boolean_flag_as_string", - "fractional_max_tokens", - "unknown_tiered_pricing_field", - ], -) -def test_schema_rejects_malformed_entries(committed_schema: dict, entry: dict): - validator = build_validator(committed_schema) - assert not validator.is_valid({"some-model": entry}) - - -def test_schema_accepts_minimal_and_unknown_optional_fields(committed_schema: dict): - validator = build_validator(committed_schema) - assert validator.is_valid({"some-model": {"litellm_provider": "openai"}}) - assert validator.is_valid({"some-model": {"litellm_provider": "openai", "brand_new_field": {"nested": True}}}) - - -def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_schema: dict): - validator = build_validator(committed_schema) - entry = { - "litellm_provider": "dashscope", - "mode": "chat", - "tiered_pricing": [ - { - "range": [0, 256000], - "input_cost_per_token": 3.25e-07, - "output_cost_per_token": 1.95e-06, - "cache_creation_input_token_cost": 4.063e-07, - "cache_read_input_token_cost": 3.25e-08, - } - ], - } - assert validator.is_valid({"some-model": entry}) - - -OFF_PEAK_ENTRY: Final = MappingProxyType( - { - "litellm_provider": "openrouter", - "mode": "chat", - "input_cost_per_token": 2e-6, - "output_cost_per_token": 8e-6, - "off_peak_pricing": { - "hours_utc": "16:30-00:30", - "windows": [{"hours_utc": ["00:30-02:00"], "weekdays": [6, "Sunday", "mon", "THURS"]}], - "weekday_timezone": "Asia/Shanghai", - "input_cost_per_token": 1e-6, - "output_cost_per_token": 4e-6, - "cache_read_input_token_cost": 1e-7, - }, - } -) - - -def test_generator_classifies_off_peak_pricing_as_a_windowed_rate_block(): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) - validator = build_validator(schema) - assert validator.is_valid({"some-model": dict(OFF_PEAK_ENTRY)}) - - -@pytest.mark.parametrize( - "block", - [ - {"hours_utc": "16:30-00:30", "input_cost_per_token": "1e-6"}, - {"hours_utc": "16:30-00:30", "input_cost_per_token": -1e-6}, - {"hours_utc": 1630, "input_cost_per_token": 1e-6}, - {"hours_utc": "16:30-00:30", "discount": 0.5}, - {"windows": [{"weekdays": [6]}], "input_cost_per_token": 1e-6}, - {"windows": [{"hours_utc": "00:30-02:00", "weekdays": [0]}], "input_cost_per_token": 1e-6}, - {"windows": [], "input_cost_per_token": 1e-6}, - {"input_cost_per_token": 1e-6}, - {"hours_utc": "16:30", "input_cost_per_token": 1e-6}, - {"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6}, - {"hours_utc": ["16:30-00:30", "4pm-midnight"], "input_cost_per_token": 1e-6}, - {"windows": [{"hours_utc": "00:30-02:00", "weekdays": ["Funday"]}], "input_cost_per_token": 1e-6}, - ], -) -def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) - validator = build_validator(schema) - assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) - - -TOKEN_COUNT_ENTRY: Final = MappingProxyType( - { - "litellm_provider": "bedrock_converse", - "mode": "chat", - "token_count_framing": 78, - "token_count_fallback_model": "us.anthropic.claude-opus-5-5", - } -) - - -@pytest.mark.parametrize("framing", [0, 78]) -def test_generator_classifies_token_count_settings(framing: int): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) - assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, "token_count_framing": framing}}) - - -@pytest.mark.parametrize( - "setting", - [ - {"token_count_framing": -1}, - {"token_count_framing": "78"}, - {"token_count_framing": 1.5}, - {"token_count_fallback_model": ""}, - {"token_count_fallback_model": 5}, - ], -) -def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) - assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) - - def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 02c90e0d97b..35fbd28eae7 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -10,13 +10,16 @@ calculations for DB-sourced models with prompt caching pricing. """ import copy +import itertools import os -from typing import Final +from collections.abc import Callable, Mapping +from typing import Final, NoReturn import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map from litellm.main import _build_custom_pricing_entry from litellm.utils import _invalidate_model_cost_lowercase_map @@ -796,6 +799,609 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): _invalidate_model_cost_lowercase_map() +@pytest.mark.parametrize("route", ["", "converse/", "invoke/"]) +def test_register_model_prices_a_geo_alias_from_its_builtin_when_the_lookup_fails( + monkeypatch: pytest.MonkeyPatch, route: str +) -> None: + from litellm.types.utils import ModelResponse, Usage + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + canonical_key: Final = "au.anthropic.claude-sonnet-5" + alias: Final = f"bedrock/{route}{canonical_key}" + canonical: Final = litellm.model_cost[canonical_key] + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model( + {alias: {"litellm_provider": "bedrock", "mode": "chat", "supports_vision": True}}, + persist_across_reloads=False, + ) + + usage: Final = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + alias_cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=alias, usage=usage), model=alias, custom_llm_provider="bedrock" + ) + canonical_cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=canonical_key, usage=usage), + model=canonical_key, + custom_llm_provider="bedrock", + ) + try: + assert alias_cost == canonical_cost > 0 + assert litellm.get_model_info(alias)["input_cost_per_token"] == canonical["input_cost_per_token"] + assert litellm.get_model_info(alias)["supports_vision"] is True + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: + geo_key: Final = "au.anthropic.lit-5853-geo-model" + geo_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "cache_read_input_token_cost": 3.3e-07, + } + sibling_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + } + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + geo_key: dict(geo_row), + "anthropic.lit-5853-geo-model": sibling_row, + }, + ) + _invalidate_model_cost_lowercase_map() + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model({geo_key: {"mode": "chat"}}, persist_across_reloads=False) + + try: + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert litellm.model_cost[geo_key][field] == geo_row[field] + assert litellm.get_model_info(geo_key)[field] == geo_row[field] + assert litellm.model_cost[geo_key]["litellm_provider"] == geo_row["litellm_provider"] + finally: + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("route", ["", "converse/", "invoke/"]) +def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once( + monkeypatch: pytest.MonkeyPatch, route: str +) -> None: + catalog_key: Final = "au.anthropic.lit-5853-geo-model" + alias: Final = f"bedrock/{route}{catalog_key}" + repriced_input_cost: Final = 9e-06 + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + catalog_key: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + }, + }, + ) + _invalidate_model_cost_lowercase_map() + real_lookup: Final = litellm.utils.get_model_info + lookups: Final = itertools.count() + + def lookup_tripped_once_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> object: + if next(lookups) == 0: + raise RuntimeError("dictionary changed size during iteration") + return real_lookup(model=model, custom_llm_provider=custom_llm_provider) + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_once_by_a_concurrent_registration) + litellm.register_model({alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False) + litellm.register_model({catalog_key: {"input_cost_per_token": repriced_input_cost}}, persist_across_reloads=False) + + try: + assert litellm.get_model_info(alias)["key"] == catalog_key + assert litellm.get_model_info(alias)["input_cost_per_token"] == repriced_input_cost + finally: + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("key_prefix", ["", "bedrock/", "bedrock/converse/"]) +def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup_fails( + monkeypatch: pytest.MonkeyPatch, key_prefix: str +) -> None: + from litellm.types.utils import ModelResponse, Usage + + catalog_key: Final = "au.anthropic.lit-5853-tiered-geo-model" + registered_key: Final = f"{key_prefix}{catalog_key}" + catalog_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "input_cost_per_token_above_200k_tokens": 6.6e-06, + "output_cost_per_token_above_200k_tokens": 2.475e-05, + } + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float: + monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)}) + _invalidate_model_cost_lowercase_map() + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup) + litellm.register_model( + { + registered_key: { + "litellm_provider": "bedrock", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + return litellm.completion_cost( + completion_response=ModelResponse( + model=registered_key, + usage=Usage(prompt_tokens=250_000, completion_tokens=1_000, total_tokens=251_000), + ), + model=registered_key, + custom_llm_provider="bedrock", + ) + + try: + assert long_prompt_cost_after_registering_with(lookup_tripped_by_a_concurrent_registration) == pytest.approx( + long_prompt_cost_after_registering_with(litellm.utils.get_model_info) + ) + finally: + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("lookup_trips", [False, True]) +def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_the_openai_row( + monkeypatch: pytest.MonkeyPatch, lookup_trips: bool +) -> None: + openai_key: Final = "lit-5853-openai-model" + azure_ai_key: Final = f"azure_ai/{openai_key}" + openai_row: Final = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 5e-06, + "output_cost_per_token": 1.5e-05, + } + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://lit-5853.openai.azure.com") + monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), openai_key: dict(openai_row)}) + monkeypatch.setattr( + litellm, "open_ai_chat_completion_models", {*litellm.open_ai_chat_completion_models, openai_key} + ) + _invalidate_model_cost_lowercase_map() + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + if lookup_trips: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model( + {azure_ai_key: {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}, + persist_across_reloads=False, + ) + + try: + assert litellm.model_cost[openai_key] == openai_row + assert litellm.cost_per_token(model=azure_ai_key, prompt_tokens=1000, completion_tokens=1000) == pytest.approx( + (1e-03, 2e-03) + ) + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monkeypatch: pytest.MonkeyPatch) -> None: + catalog_key: Final = "au.anthropic.lit-5853-reloaded-model" + deployment_key: Final = f"bedrock/{catalog_key}" + catalog_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + } + reloaded_input_cost: Final = 6.6e-06 + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)}) + _invalidate_model_cost_lowercase_map() + reloaded_cost_map: Final = { + **builtin_cost_map, + catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost}, + } + real_lookup: Final = litellm.utils.get_model_info + + def lookup_racing_a_price_data_reload(model: str, custom_llm_provider: str | None = None) -> object: + info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider) + if litellm.model_cost is not reloaded_cost_map: + adopt_model_cost_map(reloaded_cost_map) + return info + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_racing_a_price_data_reload) + litellm.register_model({deployment_key: {"litellm_provider": "bedrock"}}, persist_across_reloads=False) + + try: + assert litellm.model_cost[catalog_key]["input_cost_per_token"] == reloaded_input_cost + assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + catalog_key: Final = "au.anthropic.lit-5853-reloaded-as-written-model" + deployment_key: Final = f"bedrock/{catalog_key}" + catalog_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + } + reloaded_input_cost: Final = 6.6e-06 + registered_output_cost: Final = 2e-05 + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + reloaded_cost_map: Final = { + **builtin_cost_map, + catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost}, + } + + class CostMapAReloadReplacesAsItIsWritten(dict[str, dict[str, object]]): + def update(self, other: Mapping[str, dict[str, object]], /) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # registration writes pass one mapping + if litellm.model_cost is self: + adopt_model_cost_map(reloaded_cost_map) + super().update(other) + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + CostMapAReloadReplacesAsItIsWritten({**builtin_cost_map, catalog_key: dict(catalog_row)}), + ) + _invalidate_model_cost_lowercase_map() + + litellm.register_model( + {deployment_key: {"litellm_provider": "bedrock", "output_cost_per_token": registered_output_cost}}, + persist_across_reloads=False, + ) + + try: + assert litellm.model_cost is reloaded_cost_map + assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost + assert litellm.get_model_info(deployment_key)["output_cost_per_token"] == registered_output_cost + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: pytest.MonkeyPatch) -> None: + catalog_key: Final = "lit-5853-live-entry-model" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + catalog_key: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 2e-06}, + }, + ) + _invalidate_model_cost_lowercase_map() + walk: Final = iter(litellm.model_cost[catalog_key].items()) + first_field: Final = next(walk) + + try: + litellm.register_model( + {catalog_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06, "supports_vision": True}}, + persist_across_reloads=False, + ) + assert dict((first_field, *walk)) == { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2e-06, + } + assert litellm.get_model_info(catalog_key)["input_cost_per_token"] == 1e-06 + assert litellm.get_model_info(catalog_key)["supports_vision"] is True + finally: + _invalidate_model_cost_lowercase_map() + + +_CANONICAL_KEY: Final = "lit-5853-canonical-model" +_ALIAS_KEY: Final = "lit-5853-alias-model" + + +def _install_aliased_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + _expand_model_aliases( + { + **litellm.get_model_cost_map(url=""), + _CANONICAL_KEY: { + "aliases": [_ALIAS_KEY], + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2e-06, + "output_cost_per_token": 8e-06, + }, + } + ), + ) + _invalidate_model_cost_lowercase_map() + + +def _token_prices(name: str) -> tuple[float | None, float | None]: + info: Final = litellm.get_model_info(name) + return info["input_cost_per_token"], info["output_cost_per_token"] + + +@pytest.mark.parametrize("registered_key", [_CANONICAL_KEY, _ALIAS_KEY]) +def test_register_model_prices_every_alias_of_the_registered_entry( + monkeypatch: pytest.MonkeyPatch, registered_key: str +) -> None: + _install_aliased_cost_map(monkeypatch) + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}}, + persist_across_reloads=False, + ) + litellm.register_model( + {registered_key: {"litellm_provider": "openai", "output_cost_per_token": 4e-06}}, + persist_across_reloads=False, + ) + assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == { + _CANONICAL_KEY: (1e-06, 4e-06), + _ALIAS_KEY: (1e-06, 4e-06), + } + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_leaves_an_alias_repointed_at_its_own_entry(monkeypatch: pytest.MonkeyPatch) -> None: + _install_aliased_cost_map(monkeypatch) + litellm.model_cost[_ALIAS_KEY] = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 5e-06, + "output_cost_per_token": 6e-06, + } + + try: + litellm.register_model( + {_CANONICAL_KEY: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}}, + persist_across_reloads=False, + ) + assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == { + _CANONICAL_KEY: (1e-06, 8e-06), + _ALIAS_KEY: (5e-06, 6e-06), + } + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + builtin_key: Final = "us.anthropic.claude-sonnet-4-6" + registered_key: Final = f"bedrock/bedrock/bedrock/{builtin_key}" + builtin: Final = litellm.model_cost[builtin_key] + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + info: Final = litellm.get_model_info(registered_key) + assert info["input_cost_per_token"] == builtin["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] == builtin["output_cost_per_token"] > 0 + assert info["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"] + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_resolves_to( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.fallback_generalizations import ( + get_fallback_generalization_rules, + set_fallback_generalizations, + ) + + sibling_key: Final = "anthropic.lit-5853-flagless-model" + region_alias: Final = f"bedrock/ap-northeast-1.{sibling_key}" + sibling_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + } + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), sibling_key: sibling_row}) + _invalidate_model_cost_lowercase_map() + previous_rules: Final = list(get_fallback_generalization_rules()) + set_fallback_generalizations( + [ + { + "name": "lit-5853-flagless", + "pattern": "lit-5853-flagless", + "model_info": {"supports_system_messages": True, "supports_parallel_function_calling": True}, + } + ] + ) + + try: + unregistered: Final = litellm.get_model_info(region_alias) + litellm.register_model( + {region_alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + registered: Final = litellm.get_model_info(region_alias) + assert registered["supports_system_messages"] is unregistered["supports_system_messages"] is True + assert ( + registered["supports_parallel_function_calling"] + is unregistered["supports_parallel_function_calling"] + is True + ) + assert registered["input_cost_per_token"] == sibling_row["input_cost_per_token"] + finally: + set_fallback_generalizations(previous_rules) + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize( + ("builtin_prices", "custom_prices", "usage", "expected_cost"), + [ + pytest.param( + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + }, + {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}, + {"prompt_tokens": 250_000, "completion_tokens": 1_000}, + 250_000 * 1e-06 + 1_000 * 2e-06, + id="above-200k-tier", + ), + pytest.param( + { + "tiered_pricing": [ + {"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07, "range": [0, 256_000]}, + {"input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06, "range": [256_000, 1_000_000]}, + ] + }, + {"input_cost_per_token": 1e-06}, + {"prompt_tokens": 300_000, "completion_tokens": 1_000}, + 300_000 * 1e-06 + 1_000 * 2e-06, + id="tiered-pricing-input-only", + ), + pytest.param( + {"input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_reasoning_token": 3e-05}, + {"output_cost_per_token": 2e-06}, + {"prompt_tokens": 100, "completion_tokens": 1_000, "reasoning_tokens": 600}, + 100 * 3e-06 + 1_000 * 2e-06, + id="reasoning-rate-output-only", + ), + pytest.param( + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "off_peak_pricing": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "hours_utc": "00:00-00:00", + }, + }, + {"output_cost_per_token": 2e-06}, + {"prompt_tokens": 100, "completion_tokens": 1_000}, + 100 * 1e-06 + 1_000 * 2e-06, + id="off-peak-output-only", + ), + ], +) +def test_register_model_keeps_custom_flat_prices_over_inherited_rate_variants( + monkeypatch: pytest.MonkeyPatch, + builtin_prices: dict[str, object], + custom_prices: dict[str, float], + usage: dict[str, int], + expected_cost: float, +) -> None: + from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage + + builtin_key: Final = "anthropic.lit-5853-tiered-model-v1:0" + alias: Final = f"bedrock/ap-northeast-1.{builtin_key}" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + builtin_key: {"litellm_provider": "bedrock", "mode": "chat", **builtin_prices}, + }, + ) + _invalidate_model_cost_lowercase_map() + + try: + litellm.register_model( + {alias: {"litellm_provider": "bedrock", **custom_prices}}, + persist_across_reloads=False, + ) + response_usage: Final = Usage( + prompt_tokens=usage["prompt_tokens"], + completion_tokens=usage["completion_tokens"], + total_tokens=usage["prompt_tokens"] + usage["completion_tokens"], + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=usage.get("reasoning_tokens", 0)), + ) + cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=alias, usage=response_usage), + model=alias, + custom_llm_provider="bedrock", + ) + assert cost == pytest.approx(expected_cost) + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pytest.MonkeyPatch) -> None: + openai_key: Final = "lit-5853-openai-only-model" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + openai_key: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.25e-06, + }, + }, + ) + _invalidate_model_cost_lowercase_map() + registered_key: Final = f"bedrock/{openai_key}" + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + registered: Final = litellm.model_cost[registered_key] + assert "cache_read_input_token_cost" not in registered + assert "input_cost_per_token" not in registered + assert "output_cost_per_token" not in registered + finally: + _invalidate_model_cost_lowercase_map() + + def test_update_dictionary_merges_nested_dicts_without_aliasing(): """A nested dict must be merged copy-on-write: the pre-existing nested dict object stays untouched, and the caller's incoming nested dict is never diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d8625ff9da3..b353dc86d5f 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -9,9 +9,14 @@ should still use the built-in pricing. import asyncio import copy +import itertools import logging import os +import queue import re +import sys +import threading +from types import MappingProxyType from typing import Final from unittest.mock import Mock, patch @@ -3241,6 +3246,76 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info( assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price +def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + catalog: Final = litellm.get_model_cost_map(url="") + monkeypatch.setattr(litellm, "model_cost", dict(catalog)) + _invalidate_model_cost_lowercase_map() + backends: Final = tuple( + f"bedrock/au.anthropic.{model}" + for model in ("claude-sonnet-5", "claude-opus-5", "claude-haiku-4-5-20251001-v1:0", "claude-sonnet-4-6") + ) + canonical_prices: Final = MappingProxyType( + {backend: catalog[backend.removeprefix("bedrock/")]["input_cost_per_token"] for backend in backends} + ) + assert all(price > 0 for price in canonical_prices.values()) + router: Final = Router( + model_list=[ + { + "model_name": backend.rsplit(".", 1)[-1], + "litellm_params": {"model": backend, "aws_region_name": "ap-southeast-2"}, + "model_info": {"id": f"lit-5853-dep-{index}"}, + } + for index, backend in enumerate(backends) + ] + ) + stop: Final = threading.Event() + writer_outcome: Final[queue.SimpleQueue[int | Exception]] = queue.SimpleQueue() + + def register_until_stopped() -> None: + try: + for i in itertools.count(): + if stop.is_set(): + writer_outcome.put(i) + return + litellm.register_model( + { + f"lit-5853-writer-{i}": { + "litellm_provider": "lit-5853-writer", + "mode": "chat", + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 1e-08, + } + }, + persist_across_reloads=False, + ) + except Exception as failure: + writer_outcome.put(failure) + + writer: Final = threading.Thread(target=register_until_stopped, daemon=True) + saved_switch_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-4) + writer.start() + try: + for _ in range(10): + _simulate_price_data_reload_with_provider_sets(monkeypatch, dict(catalog)) + finally: + stop.set() + writer.join() + sys.setswitchinterval(saved_switch_interval) + + registrations_landed: Final = writer_outcome.get(timeout=5) + try: + assert isinstance(registrations_landed, int) and registrations_landed > 0, registrations_landed + assert tuple(backend for backend in backends if backend in litellm.model_cost) == () + for backend in backends: + assert litellm.get_model_info(backend)["input_cost_per_token"] == canonical_prices[backend] + assert tuple(router.get_model_names()) == tuple(backend.rsplit(".", 1)[-1] for backend in backends) + finally: + _invalidate_model_cost_lowercase_map() + + _COLLIDING_BUILTIN_KEY: Final = "baseten/zai-org/GLM-5.2" _COLLIDING_SHARED_OPENAI_KEY: Final = "openai/zai-org/GLM-5.2" _COLLIDING_MODEL_INFO: Final = { diff --git a/tests/unit/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py index 918d215c89f..6ab672d7b04 100644 --- a/tests/unit/test_router_retry_policy_update.py +++ b/tests/unit/test_router_retry_policy_update.py @@ -430,6 +430,7 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): fake_table = _FakeConfigTable() prisma_client = MagicMock() prisma_client.db.litellm_config = fake_table + prisma_client.writer_db = prisma_client.db async def _apply_router_settings(*args, **kwargs): await proxy_server.proxy_config._add_router_settings_from_db_config( diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 632ab020409..a8e90fded83 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone from pathlib import Path, PurePath -from typing import Final, cast +from typing import Final, SupportsIndex, cast from unittest import mock from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -1256,6 +1256,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_native_streaming": {"type": "boolean"}, "supports_image_size": {"type": "boolean"}, "supports_native_structured_output": {"type": "boolean"}, + "supports_bedrock_native_structured_output": {"type": "boolean"}, "use_openai_responses_path": {"type": "boolean"}, "off_peak_pricing": { "type": "object", @@ -4886,6 +4887,104 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey _invalidate_model_cost_lowercase_map() +def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import _invalidate_model_cost_lowercase_map, _rebuild_model_cost_lowercase_map + + class _KeyRegisteredMidWalk(str): + def lower(self) -> str: + litellm.model_cost.setdefault("lit-5853-registered-during-the-rebuild", {"mode": "chat"}) + return super().lower() + + class _EntryRegisteredMidWalk(dict): + def get(self, field: str, default: object = None) -> object: + litellm.model_cost.setdefault("lit-5853-registered-during-add-known-models", {"mode": "chat"}) + return super().get(field, default) + + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + litellm.model_cost[_KeyRegisteredMidWalk("LIT-5853-WALKED-LAST")] = _EntryRegisteredMidWalk(mode="chat") + try: + assert _rebuild_model_cost_lowercase_map()["lit-5853-walked-last"] == "LIT-5853-WALKED-LAST" + litellm.add_known_models(model_cost_map=litellm.model_cost) + assert "lit-5853-registered-during-the-rebuild" in litellm.model_cost + assert "lit-5853-registered-during-add-known-models" in litellm.model_cost + finally: + _invalidate_model_cost_lowercase_map() + + +class _KeyWhoseLoweringLosesTheRace(str): + """Lowering the key is the one Python-level step between the lookup map's rebuild and + its read, so another thread's registration invalidating the map lands right there.""" + + def lower(self) -> str: + from litellm.utils import _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + return super().lower() + + +def test_case_insensitive_model_cost_lookup_survives_an_invalidation_mid_lookup() -> None: + from litellm.utils import _get_model_cost_key, _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + try: + assert _get_model_cost_key(_KeyWhoseLoweringLosesTheRace("GPT-4O")) == "gpt-4o" + finally: + _invalidate_model_cost_lowercase_map() + + +def _register_a_model_mid_walk() -> None: + litellm.model_cost.setdefault("lit-5853-registered-mid-walk", {"mode": "chat"}) + + +class _KeyReadDuringARegistration(str): + def startswith( + self, + prefix: str | tuple[str, ...], + start: SupportsIndex | None = None, + end: SupportsIndex | None = None, + ) -> bool: + _register_a_model_mid_walk() + return super().startswith(prefix, start, end) + + def lower(self) -> str: + _register_a_model_mid_walk() + return super().lower() + + +class _EntryReadDuringARegistration(dict[str, object]): + def __contains__(self, key: object) -> bool: + _register_a_model_mid_walk() + return super().__contains__(key) + + +@pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5", "case-insensitive scan"]) +def test_request_path_cost_map_walks_survive_a_registration_mid_walk(monkeypatch: pytest.MonkeyPatch, walk: str) -> None: + from litellm.llms.fal_ai.cost_calculator import _keyed_rows + from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig + from litellm.llms.ollama.common_utils import OllamaModelInfo + from litellm.llms.openai.chat.gpt_5_transformation import _catalogue_declares_default_effort + from litellm.utils import _handle_new_key_with_scan, _invalidate_model_cost_lowercase_map + + walks: Final[Mapping[str, Callable[[], object]]] = { + "fal_ai": lambda: _keyed_rows("lit-5853-model", "standard"), + "fireworks_ai": FireworksAIConfig._get_fireworks_index, + "ollama": lambda: OllamaModelInfo._is_static_ollama_model("ollama/lit-5853-model"), + "gpt-5": _catalogue_declares_default_effort, + "case-insensitive scan": lambda: _handle_new_key_with_scan("lit-5853-model"), + } + monkeypatch.setattr( + litellm, + "model_cost", + {_KeyReadDuringARegistration("lit-5853-catalog-model"): _EntryReadDuringARegistration(mode="chat")}, + ) + + try: + walks[walk]() + assert "lit-5853-registered-mid-walk" in litellm.model_cost + finally: + _invalidate_model_cost_lowercase_map() + + class _JsonCapture(logging.Handler): def __init__(self): super().__init__() @@ -7090,6 +7189,25 @@ def test_get_model_info_gemini(monkeypatch): assert info.get("rpm") is not None, f"{model} does not have rpm" +def test_register_model_without_a_provider_never_inherits_another_providers_entry(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import _invalidate_model_cost_lowercase_map + + key: Final = "azure/claude-sonnet-4-6" + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + litellm.model_cost.pop(key, None) + try: + litellm.register_model( + {key: {"input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07}}, persist_across_reloads=False + ) + assert litellm.model_cost[key].get("litellm_provider") is None + info: Final = litellm.get_model_info(key) + assert info["litellm_provider"] == "azure" + assert info["input_cost_per_token"] == 1e-07 + assert info["output_cost_per_token"] == 4e-07 + finally: + _invalidate_model_cost_lowercase_map() + + @pytest.mark.parametrize( "model, litellm_params, expects_native_config", [ diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index 09c9079d89e..bfaac44a483 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -224,9 +224,20 @@ def test_regional_bedrock_alias_requires_canonical_converse_metadata( _enforce_bedrock_converse_models(model_cost, ()) +GOVCLOUD_LAGGING_CAPABILITIES: Final = frozenset({"supports_bedrock_native_structured_output"}) + + +def _govcloud_may_lag(model_key: str, capability: str) -> bool: + is_govcloud: Final = model_key.startswith("us-gov.") or any( + part.startswith("us-gov-") for part in model_key.split("/") + ) + return is_govcloud and capability in GOVCLOUD_LAGGING_CAPABILITIES + + def test_get_model_info_bedrock_models(): """ Check for drift in base model info for bedrock models and regional model info for bedrock models. + GovCloud regions may lag their commercial base on Bedrock native structured output, never lead it. """ from litellm.llms.bedrock.common_utils import BedrockModelInfo @@ -261,6 +272,9 @@ def test_get_model_info_bedrock_models(): assert ( base_model_key in v ), f"{base_model_key} is not in model cost map for {k}" + if _govcloud_may_lag(k, base_model_key): + assert base_model_value or not v[base_model_key], f"{base_model_key} is on {k} but not its base" + continue assert ( v[base_model_key] == base_model_value ), f"{base_model_key} is not equal to {base_model_value} for model {k}" @@ -297,6 +311,8 @@ def test_get_model_info_bedrock_cross_region_capability_parity(): """ Cross-region inference profiles carry litellm_provider "bedrock_converse", so the regional drift check above (which filters on "bedrock") never reaches them. + GovCloud is a separate AWS partition that gets features later, so a us-gov profile + may lack Bedrock native structured output its commercial base has, never the reverse. """ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -318,6 +334,9 @@ def test_get_model_info_bedrock_cross_region_capability_parity(): if not cap.startswith("supports_"): continue assert cap in v, f"{cap} is on {base_model_key} but missing from {k}" + if _govcloud_may_lag(k, cap): + assert base_value or not v[cap], f"{cap} is on {k} but not on {base_model_key}" + continue assert ( v[cap] == base_value ), f"{cap} is {v[cap]} on {k} but {base_value} on {base_model_key}" diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9eead41a102..423287e7c0f 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -231,9 +231,6 @@ "max-lines": { "count": 1 }, - "react-hooks/set-state-in-effect": { - "count": 1 - }, "react/no-unescaped-entities": { "count": 2 } @@ -1127,9 +1124,6 @@ "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { "react-hooks/purity": { "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/app/(dashboard)/users/_components/BulkEditUsers.tsx": { @@ -1775,11 +1769,6 @@ "count": 1 } }, - "src/components/molecules/cost_optimization_feedback_banner.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, "src/components/navbar.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index b546b2c0558..cdee86ed559 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -1,7 +1,7 @@ import { act, cleanup, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; -import type { ReactNode } from "react"; +import { useState, type ComponentProps, type ReactNode } from "react"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; import * as networking from "@/components/networking"; @@ -15,6 +15,13 @@ import type { import EntityUsage from "./EntityUsage"; import { getGlobalTopKeys, getTopAPIKeys } from "./entityUsageAggregations"; +type EntityUsageProps = ComponentProps; + +function StatefulEntityUsage(props: Omit) { + const [selectedEntities, setSelectedEntities] = useState([]); + return ; +} + const emptySpendMetrics: SpendMetrics = { spend: 0, prompt_tokens: 0, @@ -110,7 +117,7 @@ vi.mock("./TopModelView", () => ({ })); vi.mock("./TeamUserSpendCard", () => ({ - default: ({ teamIds }: { teamIds: string[] }) =>
{`team-user-spend:${teamIds.join("|")}`}
, + default: ({ teamIds }: { teamIds: readonly string[] }) =>
{`team-user-spend:${teamIds.join("|")}`}
, })); vi.mock("@/components/EntityUsageExport/EntityUsageExportModal", () => ({ @@ -549,7 +556,7 @@ describe("EntityUsage", () => { }); it("should render with tag entity type and display spend metrics", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -568,7 +575,7 @@ describe("EntityUsage", () => { }); it("should render with team entity type and call team API", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -591,12 +598,12 @@ describe("EntityUsage", () => { teamsResult([{ team_id: "team-alpha" }, { team_id: "litellm-dashboard" }, { team_id: "team-beta" }]), ); - render(); + render(); expect(await screen.findByText("team-user-spend:team-alpha|team-beta")).toBeInTheDocument(); cleanup(); mockUseTeams.mockReturnValue(teamsResult([])); - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); }); @@ -604,7 +611,7 @@ describe("EntityUsage", () => { }); it("should render with organization entity type and call organization API", async () => { - render(); + render(); await waitFor(() => { expect(mockOrganizationDailyActivityCall).toHaveBeenCalled(); @@ -619,7 +626,7 @@ describe("EntityUsage", () => { }); it("should render with customer entity type and call customer API", async () => { - render(); + render(); await waitFor(() => { expect(mockCustomerDailyActivityCall).toHaveBeenCalled(); @@ -634,7 +641,7 @@ describe("EntityUsage", () => { }); it("should render with agent entity type and call agent API", async () => { - render(); + render(); await waitFor(() => { expect(mockAgentDailyActivityCall).toHaveBeenCalled(); @@ -649,7 +656,7 @@ describe("EntityUsage", () => { }); it("should render with user entity type and call user API", async () => { - render(); + render(); await waitFor(() => { expect(mockUserDailyActivityCall).toHaveBeenCalled(); @@ -664,7 +671,7 @@ describe("EntityUsage", () => { }); it("should switch between tabs", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -688,7 +695,7 @@ describe("EntityUsage", () => { }); it("loads key pages separately from the aggregate using the current entity scope", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalledTimes(1); @@ -743,7 +750,7 @@ describe("EntityUsage", () => { ]; it.each(NON_TEAM_PANELS)("shows only the %s panel for a non-team entity type", async (tabLabel, marker) => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -769,7 +776,7 @@ describe("EntityUsage", () => { ]; it.each(TEAM_PANELS)("shows only the %s panel for the team entity type", async (tabLabel, marker) => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -800,7 +807,7 @@ describe("EntityUsage", () => { mockTagDailyActivityCall.mockResolvedValue(emptyData); - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -814,7 +821,7 @@ describe("EntityUsage", () => { }); it("should display Model Activity tab for non-agent entity types", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -824,7 +831,7 @@ describe("EntityUsage", () => { }); it("should display Request / Token Consumption tab for agent entity type", async () => { - render(); + render(); await waitFor(() => { expect(mockAgentDailyActivityCall).toHaveBeenCalled(); @@ -834,7 +841,7 @@ describe("EntityUsage", () => { }); it("should display Top Public Model Names title for non-agent entity types", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -844,7 +851,7 @@ describe("EntityUsage", () => { }); it("defaults Model Activity to public model names and toggles to litellm models", async () => { - const { container } = render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -870,7 +877,7 @@ describe("EntityUsage", () => { }); it("should display Top Agents title for agent entity type", async () => { - render(); + render(); await waitFor(() => { expect(mockAgentDailyActivityCall).toHaveBeenCalled(); @@ -885,7 +892,7 @@ describe("EntityUsage", () => { { label: "Tag 2", value: "tag-2" }, ]; - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -899,7 +906,7 @@ describe("EntityUsage", () => { it("should fallback to team_alias when entityList is provided but entity does not exist", async () => { const customEntityList = [{ label: "Tag 2", value: "tag-2" }]; - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -911,7 +918,7 @@ describe("EntityUsage", () => { }); it("should fallback to team_alias when entityList is null", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -923,7 +930,7 @@ describe("EntityUsage", () => { }); it("should still request the filter when the caller's tag scope is empty", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -933,7 +940,7 @@ describe("EntityUsage", () => { }); it("should not request the filter while the entity list is still unresolved", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -943,7 +950,7 @@ describe("EntityUsage", () => { }); it("should display Agent Activity tab for team entity type", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -953,7 +960,7 @@ describe("EntityUsage", () => { }); it("should not display Agent Activity tab for non-team entity types", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -963,7 +970,7 @@ describe("EntityUsage", () => { }); it("should display Top Agents Driving Spend card for team entity type", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -973,7 +980,7 @@ describe("EntityUsage", () => { }); it("should not display Top Agents Driving Spend card for non-team entity types", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -983,7 +990,7 @@ describe("EntityUsage", () => { }); it("should fetch agent activity data when entity type is team", async () => { - render(); + render(); await waitFor(() => { expect(mockAgentDailyActivityCall).toHaveBeenCalledWith(expect.objectContaining({ accessToken: "test-token" })); @@ -991,7 +998,7 @@ describe("EntityUsage", () => { }); it("offers per-model top keys in Model Activity but not in the agent breakdown", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -1016,7 +1023,7 @@ describe("EntityUsage", () => { resolveSpend = resolve; }), ); - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -1037,7 +1044,7 @@ describe("EntityUsage", () => { }); it("should not fetch agent activity data for non-team entity types", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -1047,7 +1054,7 @@ describe("EntityUsage", () => { }); it("should switch to Agent Activity tab for team entity type", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -1084,7 +1091,7 @@ describe("EntityUsage", () => { mockTagDailyActivityCall.mockResolvedValue(spendDataWithoutAlias); - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -1096,7 +1103,7 @@ describe("EntityUsage", () => { }); it("renders the stacked daily spend chart, the per-entity table, and the provider share bar with a $ total", async () => { - const { container } = render(); + const { container } = render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -1140,7 +1147,7 @@ describe("EntityUsage", () => { mockUserDailyActivityCall.mockResolvedValue(spendDataForUser); // entityList is null to simulate a spender missing from the paginated user list - render(); + render(); await waitFor(() => { expect(mockUserDailyActivityCall).toHaveBeenCalled(); @@ -1153,7 +1160,7 @@ describe("EntityUsage", () => { }); it("renders the provider spend table logo from the bundled provider map", async () => { - render(); + render(); const logo = await screen.findByAltText("openai logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("openai_small")); @@ -1164,7 +1171,7 @@ describe("EntityUsage", () => { ["organization", () => mockOrganizationDailyActivityCall, "Organization Spend Overview"], ["agent", () => mockAgentDailyActivityCall, "Agent Spend Overview"], ] as const)("fetches %s activity for an admin but not for an internal user", async (entityType, call, heading) => { - render(); + render(); await waitFor(() => { expect(call()).toHaveBeenCalled(); }); @@ -1172,7 +1179,7 @@ describe("EntityUsage", () => { cleanup(); call().mockClear(); - render(); + render(); expect(await screen.findByText(heading)).toBeInTheDocument(); expect(call()).not.toHaveBeenCalled(); }); @@ -1185,7 +1192,9 @@ describe("EntityUsage", () => { ["organization", () => mockOrganizationDailyActivityCall, true], ["agent", () => mockAgentDailyActivityCall, false], ] as const)("fetches %s activity for an org admin: %s", async (entityType, call, expected) => { - render(); + render( + , + ); if (expected) { await waitFor(() => { @@ -1198,7 +1207,7 @@ describe("EntityUsage", () => { }); it("keeps the team breakdown but drops its agent sub-fetch for an internal user", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -1211,7 +1220,7 @@ describe("EntityUsage", () => { }); it("keeps the tag breakdown for an internal user", async () => { - render(); + render(); await waitFor(() => { expect(mockTagDailyActivityCall).toHaveBeenCalled(); @@ -1237,7 +1246,7 @@ describe("EntityUsage", () => { }; mockTagDailyActivityCall.mockResolvedValue(spendDataUnknownProvider); - render(); + render(); await waitFor(() => { expect(screen.getAllByText("zzz-internal").length).toBeGreaterThan(0); @@ -1277,7 +1286,7 @@ describe("EntityUsage", () => { ], }); - render(); + render(); await waitFor(() => { expect(screen.getByText("top-keys:sk-abc=30.75=alice@example.com")).toBeInTheDocument(); @@ -1287,7 +1296,7 @@ describe("EntityUsage", () => { }); it("uses the aggregated team endpoint and makes a single bounded request", async () => { - render(); + render(); await waitFor(() => { expect(mockDailyActivityAggregatedCall).toHaveBeenCalledWith("team", expect.anything()); @@ -1300,7 +1309,7 @@ describe("EntityUsage", () => { }); it("does not scope the agent breakdown by the selected team ids", async () => { - render(); + render(); await waitFor(() => { expect(mockTeamDailyActivityCall).toHaveBeenCalled(); @@ -1322,7 +1331,7 @@ describe("EntityUsage", () => { }); it("does not refetch agent activity when the team selection changes", async () => { - render(); + render(); await waitFor(() => { expect(mockAgentDailyActivityCall).toHaveBeenCalled(); @@ -1345,7 +1354,7 @@ describe("EntityUsage", () => { it("surfaces a failure alert when the aggregated call fails instead of retrying other routes", async () => { mockTeamDailyActivityCall.mockRejectedValue(new Error("aggregated unavailable")); - render(); + render(); await waitFor(() => { expect(screen.getAllByText(/Fetching spend data failed/).length).toBeGreaterThan(0); @@ -1358,12 +1367,54 @@ describe("EntityUsage", () => { const userCombobox = (): HTMLElement => within(userDropdown()).getByRole("combobox"); const renderUserUsage = async () => { - render(); + render(); await waitFor(() => { expect(mockUserDailyActivityCall).toHaveBeenCalled(); }); }; + it("links the single picked user back to their Internal Users record", async () => { + render( + , + ); + + expect(await screen.findByRole("link", { name: "View user record" })).toHaveAttribute( + "href", + "/ui/users?user=user-001", + ); + }); + + it("shows no user record link until a user is picked", async () => { + await renderUserUsage(); + + expect(screen.queryByRole("link", { name: "View user record" })).not.toBeInTheDocument(); + }); + + it("reports the picked user to the owner of the selection", async () => { + const user = userEvent.setup(); + const onSelectedEntitiesChange = vi.fn(); + render( + , + ); + + await user.click(userCombobox()); + await user.click(await screen.findByText("Alice (user-001)")); + + expect(onSelectedEntitiesChange).toHaveBeenCalledWith(["user-001"]); + }); + it("offers a user filter even when the caller preloaded no user page", async () => { await renderUserUsage(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 38a3f0f84e1..8518f52e165 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -22,6 +22,7 @@ import { cn } from "@/lib/cva.config"; import React, { type ReactNode, useCallback, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import UserDropdown from "@/components/common_components/UserDropdown"; +import UserRecordLink from "../UserRecordLink"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import { UsageExportHeader } from "@/components/EntityUsageExport"; import type { EntityType } from "@/components/EntityUsageExport/types"; @@ -122,6 +123,8 @@ interface EntityUsageProps { premiumUser: boolean; dateValue: DateRangePickerValue; isOrgAdmin?: boolean; + selectedEntities: readonly string[]; + onSelectedEntitiesChange: (ids: readonly string[]) => void; } const ENTITY_CAPABILITIES: Partial> = { @@ -137,10 +140,11 @@ const EntityUsage: React.FC = ({ userRole, dateValue, isOrgAdmin = false, + selectedEntities: selectedTags, + onSelectedEntitiesChange: setSelectedTags, }) => { const { teams } = useTeams(); const teamList = useMemo(() => teams ?? [], [teams]); - const [selectedTags, setSelectedTags] = useState([]); const [modelViewType, setModelViewType] = useState("groups"); const [topKeysLimit, setTopKeysLimit] = useState(5); const [topModelsLimit, setTopModelsLimit] = useState(5); @@ -339,7 +343,10 @@ const EntityUsage: React.FC = ({ const entityFilterSlots: Partial> = { team: , user: ( - setSelectedTags(userId ? [userId] : [])} /> +
+ setSelectedTags(userId ? [userId] : [])} /> + {selectedTags.length === 1 && } +
), }; const filterSlot = entityFilterSlots[entityType]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx index cd488cbd4ea..e9de3d22110 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx @@ -25,7 +25,7 @@ interface TeamUserSpendCardProps { accessToken: string | null; startTime: Date | null; endTime: Date | null; - teamIds: string[]; + teamIds: readonly string[]; } const QUIET_HEADER = { headerClassName: "font-normal" }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index b4134157821..c37e8e3306b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -7,7 +7,9 @@ import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import { act, fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import type { UrlUpdateEvent } from "nuqs/adapters/testing"; import { renderWithProviders } from "@/../tests/test-utils"; +import { toast } from "@/lib/toast"; import { processActivityData } from "@/components/activity_metrics"; import type { Organization } from "@/components/networking"; import type { ModelActivityData } from "@/components/UsagePage/types"; @@ -16,6 +18,8 @@ import { STACKED_USAGE_PALETTE } from "@/components/shared/charts"; import { OTHER_COLOR } from "./overview/overviewData"; import UsagePage from "./UsagePageView"; +const mockEntityUsageMounted = vi.hoisted(() => vi.fn()); + // The Overview tab. The Key Activity tab stays mounted and carries its own spend-derived totals, // so assertions about the overview's request tiles are scoped here rather than to the whole page. const overview = () => within(screen.getByRole("tabpanel", { name: "Overview" })); @@ -86,14 +90,38 @@ vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ default: () =>
Top Keys
, })); -vi.mock("./EntityUsage/EntityUsage", () => ({ - default: ({ entityType, entityList }: { entityType: string; entityList: unknown }) => ( -
- Entity Usage -
- ), - EntityList: [], -})); +vi.mock("./EntityUsage/EntityUsage", async () => { + const React = await import("react"); + const EntityUsage = ({ + entityType, + entityList, + selectedEntities, + onSelectedEntitiesChange, + }: { + entityType: string; + entityList: unknown; + selectedEntities: readonly string[]; + onSelectedEntitiesChange: (ids: readonly string[]) => void; + }) => { + React.useEffect(() => { + mockEntityUsageMounted(entityType); + }, [entityType]); + return ( +
+ Entity Usage + +
+ ); + }; + return { default: EntityUsage, EntityList: [] }; +}); vi.mock("./EntityUsage/SpendByProvider", () => ({ default: () =>
Spend By Provider
, @@ -103,7 +131,8 @@ vi.mock("./EndpointUsage/EndpointUsage", () => ({ default: () =>
Endpoint Usage
, })); -vi.mock("./UsageViewSelect/UsageViewSelect", async () => { +vi.mock("./UsageViewSelect/UsageViewSelect", async (importOriginal) => { + const actual = await importOriginal(); const React = await import("react"); const UsageViewSelect = ({ value, @@ -135,14 +164,19 @@ vi.mock("./UsageViewSelect/UsageViewSelect", async () => { ); }; UsageViewSelect.displayName = "UsageViewSelect"; - return { UsageViewSelect }; + return { ...actual, UsageViewSelect }; }); -vi.mock("@/components/shared/advanced_date_picker", async () => { +vi.mock("@/components/shared/advanced_date_picker", async (importOriginal) => { + const actual = await importOriginal(); const React = await import("react"); // The button is how a test drives a range change; the real picker's own UI is // not what any test here is asserting on. - const AdvancedDatePicker = ({ onValueChange }: { onValueChange?: (value: unknown) => void }) => + const AdvancedDatePicker = ({ + onValueChange, + }: { + onValueChange?: (value: unknown, presetShortLabel: string | null) => void; + }) => React.createElement( "div", { "data-testid": "advanced-date-picker" }, @@ -152,17 +186,33 @@ vi.mock("@/components/shared/advanced_date_picker", async () => { { "data-testid": "pick-a-different-range", onClick: () => - onValueChange?.({ from: new Date("2024-01-01T00:00:00Z"), to: new Date("2024-01-08T00:00:00Z") }), + onValueChange?.( + { from: new Date(2024, 0, 1, 0, 0, 0, 0), to: new Date(2024, 0, 8, 23, 59, 59, 999) }, + null, + ), }, "pick", ), ); AdvancedDatePicker.displayName = "AdvancedDatePicker"; - return { default: AdvancedDatePicker }; + return { ...actual, default: AdvancedDatePicker }; }); vi.mock("@/components/user_agent_activity", () => ({ - default: () =>
User Agent Activity
, + default: ({ + selectedTags, + onSelectedTagsChange, + }: { + selectedTags?: readonly string[]; + onSelectedTagsChange?: (tags: readonly string[]) => void; + }) => ( +
+ User Agent Activity + +
+ ), })); vi.mock("@/components/cloudzero_export_modal", () => ({ @@ -195,6 +245,10 @@ vi.mock("@/app/(dashboard)/hooks/useIsOrgAdmin", () => ({ default: vi.fn(() => false), })); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn(() => ({ data: undefined, isPending: true, isLoading: false })), +})); + vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ useCurrentUser: vi.fn(), })); @@ -367,8 +421,10 @@ describe("UsagePage", () => { userRole: "Internal User", userRoleLabel: "Internal User", isViewOnly: false, + loginMethod: null, premiumUser: true, disabledPersonalKeyCreation: false, + passwordResetRequired: false, showSSOBanner: false, }; @@ -486,6 +542,7 @@ describe("UsagePage", () => { mockDailyActivityAggregatedCall.mockImplementation((entity: string, request: unknown) => entity === "user" ? mockUserDailyActivityAggregatedCall(request) : Promise.resolve({ results: [], metadata: {} }), ); + mockEntityUsageMounted.mockClear(); mockTagListCall.mockClear(); mockGatewayDailyActivityCall.mockClear(); mockUserDailyActivityAggregatedCall.mockResolvedValue(mockSpendData); @@ -1026,6 +1083,7 @@ describe("UsagePage", () => { act(() => { fireEvent.change(usageSelect, { target: { value: usageView } }); }); + expect(screen.queryByText("Entity Usage")).not.toBeInTheDocument(); }); @@ -1518,4 +1576,204 @@ describe("UsagePage", () => { expect(screen.getByText("Endpoint Activity")).toBeInTheDocument(); }); }); + describe("URL state", () => { + const renderAt = (searchParams: string) => { + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + const view = renderWithProviders(, { + searchParams, + onUrlUpdate, + resetUrlUpdateQueueOnMount: false, + }); + const lastQuery = (): string | undefined => onUrlUpdate.mock.lastCall?.[0].queryString; + return { ...view, onUrlUpdate, lastQuery }; + }; + + it("preselects the user filter from ?user= and links back to that user's record", async () => { + renderAt("?user=user-001"); + + const expectedRequest = { + accessToken: "test-token", + endTime: expect.any(Date), + entityIds: ["user-001"], + startTime: expect.any(Date), + }; + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith(expect.objectContaining(expectedRequest)); + }); + expect(screen.getByRole("link", { name: "View user record" })).toHaveAttribute("href", "/ui/users?user=user-001"); + }); + + it("shows no user record link while the global view is unfiltered", async () => { + renderAt(""); + + const expectedRequest = { + accessToken: "test-token", + endTime: expect.any(Date), + entityIds: null, + startTime: expect.any(Date), + }; + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith(expect.objectContaining(expectedRequest)); + }); + expect(screen.queryByRole("link", { name: "View user record" })).not.toBeInTheDocument(); + }); + + it("opens the linked entity view with its selection and writes new picks back", async () => { + const { lastQuery } = renderAt("?view=team&team=team-1"); + + const entityUsage = await screen.findByTestId("entity-usage"); + expect(entityUsage).toHaveAttribute("data-entity-type", "team"); + expect(entityUsage).toHaveAttribute("data-selected", JSON.stringify(["team-1"])); + + fireEvent.click(screen.getByRole("button", { name: "pick entities" })); + + await waitFor(() => { + expect(lastQuery()).toBe("?view=team&team=picked-1&team=picked-2"); + }); + }); + + it("writes the chosen view and drops the previous view's filters", async () => { + const { lastQuery } = renderAt("?user=user-001&tab=keys&range=30d"); + + fireEvent.change(screen.getByTestId("usage-view-select"), { target: { value: "customer" } }); + + await waitFor(() => { + expect(lastQuery()).toBe("?range=30d&view=customer"); + }); + }); + + it("restores the inner tab from the URL and writes tab changes", async () => { + const { lastQuery } = renderAt("?tab=models"); + + expect(await screen.findByRole("tab", { name: "Model Activity" })).toHaveAttribute("aria-selected", "true"); + + fireEvent.click(screen.getByRole("tab", { name: "Key Activity" })); + + await waitFor(() => { + expect(lastQuery()).toBe("?tab=keys"); + }); + }); + + it.each([ + ["users", "User Activity"], + ["errors", "Errors"], + ])("restores the %s tab from the URL", async (tab, label) => { + renderAt(`?tab=${tab}`); + + expect(await screen.findByRole("tab", { name: label })).toHaveAttribute("aria-selected", "true"); + }); + + it("renders User Agent Activity from the URL and writes selected tags back", async () => { + const { lastQuery } = renderAt("?view=user-agent-activity&tag=User-Agent%3A+Chrome%2F1.0"); + const activity = await screen.findByTestId("user-agent-activity"); + + expect(activity).toHaveAttribute("data-selected-tags", JSON.stringify(["User-Agent: Chrome/1.0"])); + fireEvent.click(screen.getByRole("button", { name: "Select agent tag" })); + + await waitFor(() => { + expect(lastQuery()).toContain("tag=User-Agent"); + }); + expect(await screen.findByTestId("user-agent-activity")).toHaveAttribute( + "data-selected-tags", + JSON.stringify(["User-Agent: Firefox/2.0"]), + ); + }); + + it("fetches the custom range in the URL and writes a newly picked one", async () => { + const { lastQuery } = renderAt("?from=2025-03-01&to=2025-03-05"); + + const expectedRequest = { + accessToken: "test-token", + endTime: new Date(2025, 2, 5, 23, 59, 59, 999), + entityIds: null, + startTime: new Date(2025, 2, 1, 0, 0, 0, 0), + }; + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith(expect.objectContaining(expectedRequest)); + }); + + fireEvent.click(screen.getByTestId("pick-a-different-range")); + + await waitFor(() => { + expect(lastQuery()).toBe("?from=2024-01-01&to=2024-01-08"); + }); + }); + + it("tells a viewer they lack access and sends them to the default view", async () => { + const warning = vi.spyOn(toast, "warning"); + mockUseAuthorized.mockReturnValue(nonAdminSession); + + const { lastQuery } = renderAt("?view=organization&org=org-1&range=30d"); + + await waitFor(() => { + expect(lastQuery()).toBe("?range=30d"); + }); + expect(warning).toHaveBeenCalledWith("You don't have access to that usage view", { + id: "usage-url-access-denied", + }); + expect(screen.queryByTestId("entity-usage")).not.toBeInTheDocument(); + }); + + it("does not render a disallowed customer view for an internal user", async () => { + mockUseAuthorized.mockReturnValue(nonAdminSession); + + renderAt("?view=customer&customer=c-1"); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + expect(mockEntityUsageMounted).not.toHaveBeenCalledWith("customer"); + expect(screen.queryByTestId("entity-usage")).not.toBeInTheDocument(); + }); + + it("does not select the admin-only Errors tab for an internal user", async () => { + mockUseAuthorized.mockReturnValue(nonAdminSession); + + const { lastQuery } = renderAt("?tab=errors"); + + expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + await waitFor(() => { + expect(lastQuery()).toBe(""); + }); + expect(screen.queryByRole("tab", { name: "Errors" })).not.toBeInTheDocument(); + }); + + it("never fetches another user's usage for a non-admin who opens their link", async () => { + const warning = vi.spyOn(toast, "warning"); + mockUseAuthorized.mockReturnValue(nonAdminSession); + + const { lastQuery } = renderAt("?user=user-001"); + + await waitFor(() => { + expect(lastQuery()).toBe(""); + }); + expect(warning).toHaveBeenCalledWith("You don't have access to that usage view", { + id: "usage-url-access-denied", + }); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalledWith( + expect.anything(), + expect.anything(), + expect.anything(), + "user-001", + ); + }); + + it("copies the current page URL as the share link", async () => { + const writeText = vi.fn().mockResolvedValue(undefined); + Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true }); + renderAt("?tab=models"); + + fireEvent.click(await screen.findByRole("button", { name: "Copy Share Link" })); + + await waitFor(() => { + expect(writeText).toHaveBeenCalledWith(window.location.href); + }); + }); + + it("shows the share link on an entity view", async () => { + renderAt("?view=team"); + + expect(await screen.findByRole("button", { name: "Copy Share Link" })).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index f91510b9f09..84cc9359ded 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -6,7 +6,8 @@ * Works at 1m+ spend logs, by querying an aggregate table instead. */ -import { Download, Search, Sparkles, X } from "lucide-react"; +import { Download, Link2, Search, Sparkles, X } from "lucide-react"; +import { useQueryStates } from "nuqs"; import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import React, { useCallback, useEffect, useMemo, useRef, useState } from "react"; @@ -21,12 +22,16 @@ import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useIsOrgAdmin from "@/app/(dashboard)/hooks/useIsOrgAdmin"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { toast } from "@/lib/toast"; import { hasCapability } from "@/utils/capabilities"; +import { copyToClipboard } from "@/utils/dataUtils"; import { all_admin_roles, internalUserRoles } from "@/utils/roles"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import CloudZeroExportModal from "@/components/cloudzero_export_modal"; import UserDropdown from "@/components/common_components/UserDropdown"; import EntityUsageExportModal from "@/components/EntityUsageExport"; +import type { EntityType } from "@/components/EntityUsageExport/types"; import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import TopUsersView from "@/components/UsagePage/components/TopUsersView"; import UserActivityPanel from "@/components/UsagePage/components/UserActivityPanel"; @@ -67,7 +72,21 @@ import ModelViewToggle, { ModelViewType } from "./ModelViewToggle"; import TopKeyView, { type TopKeyItem } from "@/components/UsagePage/components/EntityUsage/TopKeyView"; import { getGlobalTopKeys } from "./EntityUsage/entityUsageAggregations"; import UsageAIChatPanel from "./UsageAIChatPanel"; -import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; +import { allowedUsageOptions, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; +import UserRecordLink from "./UserRecordLink"; +import { + cleanUsageUrl, + dateRangeFromParams, + dateRangePatch, + entitySelectionPatch, + selectedEntitiesFromParams, + DEFAULT_USAGE_VIEW, + usageTabFromParams, + usageTabPatch, + usageUrlParsers, + usageViewFromParams, + usageViewPatch, +} from "../usageUrlState"; import { overviewTotals, rollUpBreakdown } from "./overview/overviewData"; import TopAgents from "./overview/TopAgents"; import type { AgentRow } from "./overview/agentCatalog"; @@ -89,11 +108,12 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const initialFromDate = useMemo(() => new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), []); const initialToDate = useMemo(() => new Date(), []); - // Single date state that directly triggers data fetching - const [dateValue, setDateValue] = useState({ - from: initialFromDate, - to: initialToDate, - }); + const [urlParams, setUrlParams] = useQueryStates(usageUrlParsers); + const { range, from, to } = urlParams; + const dateValue = useMemo( + () => dateRangeFromParams({ range, from, to }) ?? { from: initialFromDate, to: initialToDate }, + [range, from, to, initialFromDate, initialToDate], + ); const [fetchedTags, setFetchedTags] = useState | null>(null); // No [] default: an unresolved query must stay undefined so the customer @@ -104,36 +124,41 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const isAdmin = all_admin_roles.includes(userRole || ""); const canViewTagUsage = isAdmin || internalUserRoles.includes(userRole || ""); const isOrgAdmin = useIsOrgAdmin(); - const canViewOrganizationUsage = hasCapability(userRole, "viewOrganizationUsage", isOrgAdmin); const canViewAgentUsage = hasCapability(userRole, "viewAgentUsage"); + const allowedViews = useMemo( + () => allowedUsageOptions(userRole, canViewTagUsage, isOrgAdmin), + [userRole, canViewTagUsage, isOrgAdmin], + ); - // For admins: null means global view (all users), a string means filter by that user - // For non-admins: always set to their own user ID - const [selectedUserId, setSelectedUserId] = useState(isAdmin ? null : userID || null); const [modelViewType, setModelViewType] = useState("groups"); const [modelQuery, setModelQuery] = useState(""); const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false); const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false); const [isAiChatOpen, setIsAiChatOpen] = useState(false); - const [selectedUsageView, setUsageView] = useState("global"); - // Org-admin membership is read from the server, so unlike the other usage - // views this one can be revoked while the page is open. Derive the view in - // render rather than storing it, so the fallback lands on the same paint and - // the selector never holds a value it no longer offers. - const usageView: UsageOption = - selectedUsageView === "organization" && !canViewOrganizationUsage ? "global" : selectedUsageView; + const selectedUsageView = usageViewFromParams(urlParams); + const selectedTab = usageTabFromParams(urlParams, isAdmin); + const usageView = allowedViews.includes(selectedUsageView) ? selectedUsageView : DEFAULT_USAGE_VIEW; const [showCredentialBanner, setShowCredentialBanner] = useState(true); const [topKeysLimit, setTopKeysLimit] = useState(5); - // Sync selectedUserId when auth state settles (isAdmin/userID may be null on initial render) - useEffect(() => { - if (!isAdmin && userID) { - setSelectedUserId(userID); - } - }, [isAdmin, userID]); - // For non-admins or "my-usage" view, always pass their own user_id - const effectiveUserId = usageView === "my-usage" || !isAdmin ? userID || null : selectedUserId; + const effectiveUserId = usageView === "my-usage" || !isAdmin ? userID || null : urlParams.user; + + const { isLoading: organizationsLoading } = useOrganizations(); + const accessSettled = userRole !== null && !organizationsLoading; + useEffect(() => { + if (!accessSettled) return; + const cleanup = cleanUsageUrl(urlParams, { allowedViews, isAdmin }); + if (!cleanup) return; + if (cleanup.deniedAccess) + toast.warning("You don't have access to that usage view", { id: "usage-url-access-denied" }); + void setUrlParams(cleanup.patch); + }, [accessSettled, urlParams, allowedViews, isAdmin, setUrlParams]); + + const entitySelection = (entityType: EntityType) => ({ + selectedEntities: selectedEntitiesFromParams(urlParams, entityType), + onSelectedEntitiesChange: (ids: readonly string[]) => void setUrlParams(entitySelectionPatch(entityType, ids)), + }); const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]); const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]); @@ -211,11 +236,15 @@ const UsagePage: React.FC = ({ teams, organizations }) => { deps: [accessToken, startTime, endTime, showTopAgents], }); const tagDaily = useMemo(() => toDailyData(tagDailyRaw), [tagDailyRaw]); - const [agentActivityTags, setAgentActivityTags] = useState(undefined); - const openAgentActivity = useCallback((agent: AgentRow) => { - setAgentActivityTags(agent.tags); - setUsageView("user-agent-activity"); - }, []); + const openAgentActivity = useCallback( + (agent: AgentRow) => { + void setUrlParams({ + ...usageViewPatch("user-agent-activity"), + ...entitySelectionPatch("tag", agent.tags), + }); + }, + [setUrlParams], + ); // Gateway request counts (SGR). Admin-only: the source table is // deployment-wide, so a non-admin must not see it. @@ -277,9 +306,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { [userSpendData], ); - const handleDateChange = useCallback((newValue: DateRangePickerValue) => { - setDateValue(newValue); - }, []); + const handleDateChange = useCallback( + (newValue: DateRangePickerValue, presetShortLabel: string | null) => + void setUrlParams(dateRangePatch(newValue, presetShortLabel)), + [setUrlParams], + ); const totals = useMemo( () => overviewTotals(userSpendData.metadata, gatewayActivity), @@ -360,18 +391,27 @@ const UsagePage: React.FC = ({ teams, organizations }) => {
setUsageView(value)} + onChange={(value) => void setUrlParams(usageViewPatch(value))} userRole={userRole} canViewTagUsage={canViewTagUsage} isOrgAdmin={isOrgAdmin} />
{isAdmin && usageView === "global" && ( -
- +
+
+ void setUrlParams({ user })} /> +
+ {urlParams.user && }
)} - +
+ + +
{aggregatedFailed && ( @@ -387,7 +427,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { {/* Your Usage / Global Usage Panel */} {(usageView === "global" || usageView === "my-usage") && ( <> - + void setUrlParams(usageTabPatch(tab))} + className="gap-3" + >
@@ -555,10 +599,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { )} {/* Organization Usage Panel */} - {usageView === "organization" && canViewOrganizationUsage && ( + {usageView === "organization" && ( = ({ teams, organizations }) => { = ({ teams, organizations }) => { = ({ teams, organizations }) => { = ({ teams, organizations }) => { = ({ teams, organizations }) => { = ({ teams, organizations }) => { {/* User Agent Activity Panel */} {usageView === "user-agent-activity" && ( void setUrlParams(entitySelectionPatch("tag", tags))} /> )}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx index 42ae8407d3f..70cfcbc7950 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx @@ -104,6 +104,27 @@ const OPTIONS: OptionConfig[] = [ adminOnly: true, }, ]; +const visibleOptions = (userRole: string | null, canViewTagUsage: boolean, isOrgAdmin: boolean): OptionConfig[] => { + const isAdmin = all_admin_roles.includes(userRole ?? ""); + return OPTIONS.filter((option) => { + if (option.capability) { + return hasCapability(userRole, option.capability, isOrgAdmin); + } + if (option.value === "tag" && canViewTagUsage) { + return true; + } + return !(option.adminOnly && !isAdmin); + }); +}; + +export const USAGE_OPTIONS: readonly UsageOption[] = OPTIONS.map((option) => option.value); + +export const allowedUsageOptions = ( + userRole: string | null, + canViewTagUsage: boolean, + isOrgAdmin: boolean, +): readonly UsageOption[] => visibleOptions(userRole, canViewTagUsage, isOrgAdmin).map((option) => option.value); + export const UsageViewSelect: React.FC = ({ value, onChange, @@ -116,18 +137,7 @@ export const UsageViewSelect: React.FC = ({ }) => { const isAdmin = all_admin_roles.includes(userRole ?? ""); const getFilteredOptions = () => { - return OPTIONS.filter((option) => { - if (option.capability) { - return hasCapability(userRole, option.capability, isOrgAdmin); - } - if (option.value === "tag" && canViewTagUsage) { - return true; - } - if (option.adminOnly && !isAdmin) { - return false; - } - return true; - }).map((option) => { + return visibleOptions(userRole, canViewTagUsage, isOrgAdmin).map((option) => { let label = option.label; let desc = option.description; if (option.showForAdmin && option.showForNonAdmin) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UserRecordLink.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UserRecordLink.tsx new file mode 100644 index 00000000000..492a0f75686 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UserRecordLink.tsx @@ -0,0 +1,16 @@ +import { SquareArrowOutUpRight } from "lucide-react"; +import Link from "next/link"; + +import { userRecordHref } from "../usageUrlState"; + +export default function UserRecordLink({ userId }: { userId: string }) { + return ( + + View user record + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/usageUrlState.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/usageUrlState.test.ts new file mode 100644 index 00000000000..dbb5c44cc35 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/usageUrlState.test.ts @@ -0,0 +1,267 @@ +import moment from "moment"; +import { createLoader, createSerializer } from "nuqs"; +import { describe, expect, it } from "vitest"; + +import type { UsageOption } from "./components/UsageViewSelect/UsageViewSelect"; +import { USAGE_OPTIONS } from "./components/UsageViewSelect/UsageViewSelect"; +import { + cleanUsageUrl, + dateRangeFromParams, + dateRangePatch, + entitySelectionPatch, + selectedEntitiesFromParams, + usageHrefForUser, + usageTabFromParams, + usageTabPatch, + usageUrlParsers, + usageViewFromParams, + usageViewPatch, + type UsageUrlParams, + type UsageUrlPatch, +} from "./usageUrlState"; + +const load = createLoader(usageUrlParsers); +const serialize = createSerializer(usageUrlParsers); + +const applyPatch = (query: string, patch: UsageUrlPatch): string => serialize(query, patch); + +const withTimezone = (timezone: string, callback: () => T): T => { + const originalTimezone = process.env.TZ; + process.env.TZ = timezone; + try { + return callback(); + } finally { + if (originalTimezone === undefined) { + delete process.env.TZ; + } else { + process.env.TZ = originalTimezone; + } + } +}; + +const adminAccess = { allowedViews: USAGE_OPTIONS, isAdmin: true }; +const internalUserAccess = { + allowedViews: ["global", "tag"] as readonly UsageOption[], + isAdmin: false, +}; + +const clean = (query: string, access = adminAccess): { query: string; deniedAccess: boolean } | null => { + const cleanup = cleanUsageUrl(load(query), access); + return cleanup && { query: applyPatch(query, cleanup.patch), deniedAccess: cleanup.deniedAccess }; +}; + +describe("usage view and tab", () => { + it("reads a known view and tab from the URL", () => { + const params = load("?view=team&tab=models"); + + expect(usageViewFromParams(params)).toBe("team"); + expect(usageTabFromParams(params)).toBe("models"); + }); + + it("falls back to the global view and cost tab when the URL names neither", () => { + const params = load("?view=bogus&tab=bogus"); + + expect(usageViewFromParams(params)).toBe("global"); + expect(usageTabFromParams(params)).toBe("cost"); + }); + + it("switching views drops the previous view's filters and tab but keeps the date range", () => { + expect(applyPatch("?user=u1&tab=keys&range=30d", usageViewPatch("team"))).toBe("?range=30d&view=team"); + expect(applyPatch("?view=team&team=t1", usageViewPatch("global"))).toBe(""); + }); + + it("round-trips the User Agent Activity view", () => { + const query = applyPatch("", usageViewPatch("user-agent-activity")); + + expect(query).toBe("?view=user-agent-activity"); + expect(usageViewFromParams(load(query))).toBe("user-agent-activity"); + }); + + it.each(["users", "errors"] as const)("round-trips the %s tab", (tab) => { + const query = applyPatch("", usageTabPatch(tab)); + + expect(query).toBe(`?tab=${tab}`); + expect(usageTabFromParams(load(query))).toBe(tab); + }); + + it("falls back from the admin-only Errors tab for an internal user", () => { + expect(usageTabFromParams(load("?tab=errors"), false)).toBe("cost"); + expect(clean("?tab=errors", internalUserAccess)).toEqual({ query: "", deniedAccess: false }); + }); + + it("keeps the default tab out of the URL", () => { + expect(applyPatch("?tab=models", usageTabPatch("cost"))).toBe(""); + expect(applyPatch("", usageTabPatch("mcp"))).toBe("?tab=mcp"); + }); +}); + +describe("entity selection", () => { + it("stores teams as repeated keys so ids with commas survive", () => { + const query = applyPatch("?view=tag", entitySelectionPatch("tag", ["Credential: a,b", "prod"])); + + expect(query).toBe("?view=tag&tag=Credential:+a,b&tag=prod"); + expect(selectedEntitiesFromParams(load(query), "tag")).toEqual(["Credential: a,b", "prod"]); + }); + + it("maps the organization view to the org key", () => { + const query = applyPatch("?view=organization", entitySelectionPatch("organization", ["o1"])); + + expect(query).toBe("?view=organization&org=o1"); + expect(selectedEntitiesFromParams(load(query), "organization")).toEqual(["o1"]); + }); + + it("stores the user view's selection in the single user key", () => { + const query = applyPatch("?view=user", entitySelectionPatch("user", ["u1"])); + + expect(query).toBe("?view=user&user=u1"); + expect(selectedEntitiesFromParams(load(query), "user")).toEqual(["u1"]); + }); + + it("keeps User Agent Activity tag filters in the URL", () => { + const query = applyPatch("?view=user-agent-activity", entitySelectionPatch("tag", ["User-Agent: Chrome/1.0"])); + + expect(query).toBe("?view=user-agent-activity&tag=User-Agent:+Chrome/1.0"); + expect(clean(query)).toBeNull(); + expect(selectedEntitiesFromParams(load(query), "tag")).toEqual(["User-Agent: Chrome/1.0"]); + }); + + it("removes the key once the selection is cleared", () => { + expect(applyPatch("?view=team&team=t1", entitySelectionPatch("team", []))).toBe("?view=team"); + expect(applyPatch("?view=user&user=u1", entitySelectionPatch("user", []))).toBe("?view=user"); + }); +}); + +describe("date range", () => { + it("stores a preset as a rolling range key", () => { + const lastThirtyDays = { from: moment().subtract(30, "days").startOf("day").toDate(), to: new Date() }; + + expect(applyPatch("?from=a&to=b", dateRangePatch(lastThirtyDays, "30d"))).toBe("?range=30d"); + }); + + it("stores a custom range matching today as explicit dates", () => { + const today = moment(); + const value = { from: today.clone().startOf("day").toDate(), to: today.clone().endOf("day").toDate() }; + + expect(dateRangePatch(value, null)).toEqual({ + range: null, + from: today.format("YYYY-MM-DD"), + to: today.format("YYYY-MM-DD"), + }); + }); + + it("stores an unknown preset label as explicit dates", () => { + const today = moment(); + const value = { from: today.clone().startOf("day").toDate(), to: today.clone().endOf("day").toDate() }; + + expect(dateRangePatch(value, "not-a-preset")).toEqual({ + range: null, + from: today.format("YYYY-MM-DD"), + to: today.format("YYYY-MM-DD"), + }); + }); + + it("keeps the custom calendar day when opened in a different timezone", () => { + const query = withTimezone("Australia/Sydney", () => { + const value = { + from: new Date(2026, 0, 5, 0, 0, 0, 0), + to: new Date(2026, 0, 5, 23, 59, 59, 999), + }; + + expect(value.from.getTimezoneOffset()).toBe(-660); + expect(moment(value.from).format("YYYY-MM-DD")).toBe("2026-01-05"); + expect(moment(value.to).format("YYYY-MM-DD")).toBe("2026-01-05"); + + const serialized = applyPatch("", dateRangePatch(value, null)); + expect(serialized).toBe("?from=2026-01-05&to=2026-01-05"); + return serialized; + }); + + withTimezone("America/New_York", () => { + expect(new Date(2026, 0, 5).getTimezoneOffset()).toBe(300); + const range = dateRangeFromParams(load(query)); + expect(range).toBeDefined(); + if (!range) return; + + expect(moment(range.from).format("YYYY-MM-DD")).toBe("2026-01-05"); + expect(moment(range.to).format("YYYY-MM-DD")).toBe("2026-01-05"); + }); + }); + + it("re-resolves a rolling range against the current day", () => { + const range = dateRangeFromParams(load("?range=7d")); + + expect(moment(range?.from).isSame(moment().subtract(7, "days"), "day")).toBe(true); + expect(moment(range?.to).isSame(moment(), "day")).toBe(true); + }); + + it("round-trips a custom range as whole local calendar days", () => { + const custom = { from: new Date(2026, 0, 5, 8, 0, 0, 0), to: new Date(2026, 0, 9, 20, 30, 0, 0) }; + const query = applyPatch("?range=7d", dateRangePatch(custom, null)); + + expect(query).toBe("?from=2026-01-05&to=2026-01-09"); + expect(dateRangeFromParams(load(query))).toEqual({ + from: new Date(2026, 0, 5, 0, 0, 0, 0), + to: new Date(2026, 0, 9, 23, 59, 59, 999), + }); + }); + + it.each([ + ["an unknown preset", "?range=90d"], + ["an ISO timestamp", "?from=2026-01-04T13:00:00.000Z&to=2026-01-05"], + ["an impossible date", "?from=2026-02-30&to=2026-03-01"], + ["a malformed date", "?from=yesterday&to=2026-01-09"], + ["a range that ends before it starts", "?from=2026-01-09&to=2026-01-05"], + ["only one end of a custom range", "?from=2026-01-05"], + ])("ignores %s", (_label, query) => { + expect(dateRangeFromParams(load(query))).toBeUndefined(); + }); +}); + +describe("cleanUsageUrl", () => { + it("leaves a valid URL alone", () => { + expect(clean("?view=team&team=t1&team=t2&range=30d")).toBeNull(); + expect(clean("?user=u1&tab=keys&from=2026-01-05&to=2026-01-09")).toBeNull(); + expect(clean("?view=tag&tag=prod", internalUserAccess)).toBeNull(); + }); + + it("silently drops params that name nothing real", () => { + expect(clean("?view=bogus&tab=bogus&range=90d")).toEqual({ query: "", deniedAccess: false }); + expect(clean("?from=yesterday&to=today&team=t1")).toEqual({ query: "", deniedAccess: false }); + }); + + it("drops a tab on a view that has no tabs", () => { + expect(clean("?view=team&tab=models")).toEqual({ query: "?view=team", deniedAccess: false }); + }); + + it("drops filters that belong to a different view", () => { + expect(clean("?view=team&team=t1&org=o1&user=u1")).toEqual({ query: "?view=team&team=t1", deniedAccess: false }); + expect(clean("?view=my-usage&user=u1")).toEqual({ query: "?view=my-usage", deniedAccess: false }); + }); + + it("prefers the preset when a URL carries both a preset and a custom range", () => { + expect(clean("?range=7d&from=2026-01-05&to=2026-01-09")).toEqual({ + query: "?range=7d", + deniedAccess: false, + }); + }); + + it("sends a viewer to the default view, keeping the dates, when the linked view is not theirs", () => { + expect(clean("?view=team&team=t1&range=30d", internalUserAccess)).toEqual({ + query: "?range=30d", + deniedAccess: true, + }); + }); + + it("denies another user's usage to a non-admin", () => { + expect(clean("?user=someone-else&tab=models", internalUserAccess)).toEqual({ query: "", deniedAccess: true }); + }); +}); + +describe("cross-page links", () => { + it("encodes the user id so it survives the round trip into the usage URL", () => { + const href = usageHrefForUser("alice+ops@example.com"); + const params: UsageUrlParams = load(href.slice(href.indexOf("?"))); + + expect(params.user).toBe("alice+ops@example.com"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/usageUrlState.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/usageUrlState.ts new file mode 100644 index 00000000000..a1cd04b723a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/usageUrlState.ts @@ -0,0 +1,180 @@ +import moment from "moment"; +import { parseAsNativeArrayOf, parseAsString, type inferParserType } from "nuqs"; + +import type { EntityType } from "@/components/EntityUsageExport/types"; +import { relativeTimeOptions } from "@/components/shared/advanced_date_picker"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; +import { uiHref } from "@/utils/uiHref"; +import { USAGE_OPTIONS, type UsageOption } from "./components/UsageViewSelect/UsageViewSelect"; + +export const DEFAULT_USAGE_VIEW: UsageOption = "global"; +export const USAGE_TABS = ["cost", "models", "keys", "users", "mcp", "endpoints", "errors"] as const; +export type UsageTab = (typeof USAGE_TABS)[number]; +export const DEFAULT_USAGE_TAB: UsageTab = "cost"; + +const entityIds = parseAsNativeArrayOf(parseAsString); + +export const usageUrlParsers = { + view: parseAsString, + tab: parseAsString, + range: parseAsString, + from: parseAsString, + to: parseAsString, + user: parseAsString, + team: entityIds, + org: entityIds, + customer: entityIds, + tag: entityIds, + agent: entityIds, +}; + +export type UsageUrlParams = inferParserType; +export type UsageUrlPatch = { readonly [K in keyof UsageUrlParams]?: UsageUrlParams[K] | null }; + +type MultiEntityType = Exclude; +type EntityKey = "user" | (typeof MULTI_ENTITY_KEYS)[MultiEntityType]; + +const MULTI_ENTITY_KEYS = { + organization: "org", + team: "team", + customer: "customer", + tag: "tag", + agent: "agent", +} as const satisfies Record; + +const VIEW_ENTITY_KEYS: Readonly> = { + global: ["user"], + "my-usage": [], + organization: ["org"], + team: ["team"], + customer: ["customer"], + tag: ["tag"], + agent: ["agent"], + user: ["user"], + "user-agent-activity": ["tag"], +}; + +const TABBED_VIEWS: readonly UsageOption[] = ["global", "my-usage"]; + +const RESET_TO_DEFAULT_VIEW: UsageUrlPatch = { + view: null, + tab: null, + user: null, + team: null, + org: null, + customer: null, + tag: null, + agent: null, +}; + +const isUsageOption = (value: string | null): value is UsageOption => USAGE_OPTIONS.some((option) => option === value); + +const isUsageTab = (value: string | null): value is UsageTab => USAGE_TABS.some((tab) => tab === value); + +const hasEntityValue = (params: UsageUrlParams, key: EntityKey): boolean => + key === "user" ? params.user !== null : params[key].length > 0; + +export const usageViewFromParams = (params: UsageUrlParams): UsageOption => + isUsageOption(params.view) ? params.view : DEFAULT_USAGE_VIEW; + +export const usageTabFromParams = (params: UsageUrlParams, isAdmin = true): UsageTab => + isUsageTab(params.tab) && (params.tab !== "errors" || isAdmin) ? params.tab : DEFAULT_USAGE_TAB; + +export const usageViewPatch = (view: UsageOption): UsageUrlPatch => ({ + ...RESET_TO_DEFAULT_VIEW, + view: view === DEFAULT_USAGE_VIEW ? null : view, +}); + +export const usageTabPatch = (tab: string): UsageUrlPatch => ({ + tab: isUsageTab(tab) && tab !== DEFAULT_USAGE_TAB ? tab : null, +}); + +export const selectedEntitiesFromParams = (params: UsageUrlParams, entityType: EntityType): readonly string[] => { + if (entityType === "user") return params.user === null ? [] : [params.user]; + return params[MULTI_ENTITY_KEYS[entityType]]; +}; + +export const entitySelectionPatch = (entityType: EntityType, ids: readonly string[]): UsageUrlPatch => { + if (entityType === "user") return { user: ids[0] ?? null }; + return { [MULTI_ENTITY_KEYS[entityType]]: ids.length > 0 ? [...ids] : null }; +}; + +const customRangeFromParams = (from: string, to: string): DateRangePickerValue | undefined => { + const fromDate = moment(from, "YYYY-MM-DD", true); + const toDate = moment(to, "YYYY-MM-DD", true); + if (!fromDate.isValid() || !toDate.isValid() || fromDate.isAfter(toDate, "day")) return undefined; + return { from: fromDate.startOf("day").toDate(), to: toDate.endOf("day").toDate() }; +}; + +type DateParams = Pick; + +export const dateRangeFromParams = (params: DateParams): DateRangePickerValue | undefined => { + if (params.range !== null) { + return relativeTimeOptions.find((option) => option.shortLabel === params.range)?.getValue(); + } + if (params.from === null || params.to === null) return undefined; + return customRangeFromParams(params.from, params.to); +}; + +export const dateRangePatch = (value: DateRangePickerValue, presetShortLabel: string | null): UsageUrlPatch => { + if (presetShortLabel !== null && relativeTimeOptions.some((option) => option.shortLabel === presetShortLabel)) { + return { range: presetShortLabel, from: null, to: null }; + } + if (!value.from || !value.to) return { range: null, from: null, to: null }; + return { + range: null, + from: moment(value.from).format("YYYY-MM-DD"), + to: moment(value.to).format("YYYY-MM-DD"), + }; +}; + +const invalidDatePatch = (params: DateParams): UsageUrlPatch => { + const hasCustomRange = params.from !== null || params.to !== null; + if (params.range !== null) { + const rangeIsPreset = relativeTimeOptions.some((option) => option.shortLabel === params.range); + return { + ...(rangeIsPreset ? {} : { range: null }), + ...(hasCustomRange ? { from: null, to: null } : {}), + }; + } + if (!hasCustomRange || dateRangeFromParams(params)) return {}; + return { from: null, to: null }; +}; + +export interface UsageUrlAccess { + readonly allowedViews: readonly UsageOption[]; + readonly isAdmin: boolean; +} + +export interface UsageUrlCleanup { + readonly patch: UsageUrlPatch; + readonly deniedAccess: boolean; +} + +export const cleanUsageUrl = (params: UsageUrlParams, access: UsageUrlAccess): UsageUrlCleanup | null => { + const view = usageViewFromParams(params); + const deniedView = !access.allowedViews.includes(view); + const deniedUserFilter = view === "global" && params.user !== null && !access.isAdmin; + if (deniedView || deniedUserFilter) { + return { patch: { ...RESET_TO_DEFAULT_VIEW, ...invalidDatePatch(params) }, deniedAccess: true }; + } + + const viewKeys = VIEW_ENTITY_KEYS[view]; + const strayEntityKeys = (["user", "team", "org", "customer", "tag", "agent"] as const).filter( + (key) => !viewKeys.includes(key) && hasEntityValue(params, key), + ); + const invalidTab = + params.tab !== null && + (!TABBED_VIEWS.includes(view) || !isUsageTab(params.tab) || (!access.isAdmin && params.tab === "errors")); + const patch: UsageUrlPatch = { + ...(params.view !== null && !isUsageOption(params.view) ? { view: null } : {}), + ...(invalidTab ? { tab: null } : {}), + ...Object.fromEntries(strayEntityKeys.map((key) => [key, null])), + ...invalidDatePatch(params), + }; + return Object.keys(patch).length > 0 ? { patch, deniedAccess: false } : null; +}; + +export const usageHrefForUser = (userId: string): string => uiHref(`usage?${new URLSearchParams({ user: userId })}`); + +export const userRecordHref = (userId: string): string => uiHref(`users?${new URLSearchParams({ user: userId })}`); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 5b98ca82c16..1dbab3b1803 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -17,7 +17,7 @@ import { import OnboardingModal, { InvitationLink } from "@/components/onboarding_link"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; -import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; +import { all_admin_roles, isAdminRole, isProxyAdminRole } from "@/utils/roles"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { @@ -298,6 +298,7 @@ const ViewUserDashboard: React.FC = ({ isLoading={userListQuery.isLoading || userListQuery.isPlaceholderData} possibleUIRoles={possibleUIRoles} teams={teams} + canViewUserUsage={all_admin_roles.includes(userRole ?? "")} sorting={sorting} onSortingChange={handleSortingChange} pagination={pagination} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx index 16911946d34..b9ed37a2152 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx @@ -36,6 +36,7 @@ interface HarnessOverrides { rowCount?: number; isLoading?: boolean; selectionEnabled?: boolean; + canViewUserUsage?: boolean; onUserClick?: (userId: string, openInEditMode?: boolean) => void; onDeleteUser?: (user: UserInfo) => void; onResetPassword?: (userId: string) => void; @@ -51,6 +52,7 @@ function Harness({ rowCount = 1, isLoading = false, selectionEnabled = false, + canViewUserUsage = true, onUserClick = vi.fn(), onDeleteUser = vi.fn(), onResetPassword = vi.fn(), @@ -74,6 +76,7 @@ function Harness({ isLoading={isLoading} possibleUIRoles={possibleUIRoles} teams={[]} + canViewUserUsage={canViewUserUsage} sorting={sorting} onSortingChange={(updater) => { setSorting(updater); @@ -178,6 +181,27 @@ describe("UsersTable", () => { expect(onUserClick).toHaveBeenCalledWith("user-1", true); }); + it("links the row menu to that user's usage when allowed", async () => { + const user = userEvent.setup(); + render(); + + await openRowMenu(user, "user-1"); + + expect(await screen.findByRole("menuitem", { name: "View usage" })).toHaveAttribute( + "href", + "/ui/usage?user=user-1", + ); + }); + + it("hides the usage link when the viewer cannot view usage", async () => { + const user = userEvent.setup(); + render(); + + await openRowMenu(user, "user-1"); + + expect(screen.queryByRole("menuitem", { name: "View usage" })).not.toBeInTheDocument(); + }); + it("delegates delete and reset-password from the row menu", async () => { const user = userEvent.setup(); const onDeleteUser = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx index 8fa76a64af6..3e266ef482d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx @@ -33,6 +33,7 @@ interface UsersTableProps { isLoading: boolean; possibleUIRoles: Record> | null; teams: UsersTableTeamOption[] | null; + canViewUserUsage: boolean; sorting: SortingState; onSortingChange: OnChangeFn; pagination: PaginationState; @@ -74,6 +75,7 @@ export function UsersTable({ isLoading, possibleUIRoles, teams, + canViewUserUsage, sorting, onSortingChange, pagination, @@ -95,12 +97,13 @@ export function UsersTable({ const columnDeps = { possibleUIRoles, includeSelection: selectionEnabled, + canViewUserUsage, onUserClick, onDeleteUser, onResetPassword, }; return getUsersTableColumns(columnDeps); - }, [possibleUIRoles, selectionEnabled, onUserClick, onDeleteUser, onResetPassword]); + }, [possibleUIRoles, selectionEnabled, canViewUserUsage, onUserClick, onDeleteUser, onResetPassword]); const roleOptions = useMemo( () => diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx index d14e2e0f4b1..ad0cdbbd9c1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx @@ -1,7 +1,8 @@ "use client"; import { ColumnDef } from "@tanstack/react-table"; -import { Copy, Info, KeyRound, MoreHorizontal, Pencil, Trash2 } from "lucide-react"; +import { ChartColumn, Copy, Info, KeyRound, MoreHorizontal, Pencil, Trash2 } from "lucide-react"; +import Link from "next/link"; import { UserInfo } from "@/components/networking"; import { createSelectionColumn, DataTableSortHeader } from "@/components/shared/DataTable"; @@ -16,6 +17,7 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; +import { usageHrefForUser } from "@/app/(dashboard)/usage/_components/usageUrlState"; import { copyToClipboard } from "@/utils/dataUtils"; const SSO_ID_HINT = @@ -29,12 +31,13 @@ function isScimInactive(user: UserInfo): boolean { interface UserRowActionsProps { user: UserInfo; + canViewUserUsage: boolean; onUserClick: (userId: string, openInEditMode?: boolean) => void; onDeleteUser: (user: UserInfo) => void; onResetPassword: (userId: string) => void; } -function UserRowActions({ user, onUserClick, onDeleteUser, onResetPassword }: UserRowActionsProps) { +function UserRowActions({ user, canViewUserUsage, onUserClick, onDeleteUser, onResetPassword }: UserRowActionsProps) { return ( Copy user ID + {canViewUserUsage && ( + } + data-testid="user-action-view-usage" + > + + View usage + + )} onDeleteUser(user)} data-testid="user-action-delete"> @@ -73,6 +85,7 @@ function UserRowActions({ user, onUserClick, onDeleteUser, onResetPassword }: Us export interface UsersTableColumnsDeps { possibleUIRoles: Record> | null; includeSelection: boolean; + canViewUserUsage: boolean; onUserClick: (userId: string, openInEditMode?: boolean) => void; onDeleteUser: (user: UserInfo) => void; onResetPassword: (userId: string) => void; @@ -81,6 +94,7 @@ export interface UsersTableColumnsDeps { export const getUsersTableColumns = ({ possibleUIRoles, includeSelection, + canViewUserUsage, onUserClick, onDeleteUser, onResetPassword, @@ -257,6 +271,7 @@ export const getUsersTableColumns = ({
= ({ dateRange, selectedFilters }) => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index f6940286051..58373eb08e4 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -27,7 +27,7 @@ interface UsageExportHeaderProps { showFilters?: boolean; filterLabel?: string; filterPlaceholder?: string; - selectedFilters?: string[]; + selectedFilters?: readonly string[]; onFiltersChange?: (filters: string[]) => void; filterOptions?: Array<{ label: string; value: string }>; filterSlot?: React.ReactNode; @@ -81,7 +81,7 @@ const UsageExportHeader: React.FC = ({ multiple disabled={isFilterDisabled} items={optionValues} - value={selectedFilters} + value={[...selectedFilters]} onValueChange={(next: string[]) => onFiltersChange?.(next)} > } className="w-full"> diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 189a14392ec..c2b42e58c80 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -11,7 +11,7 @@ export interface EntityUsageExportModalProps { entityType: EntityType; onExport: (exportType: ExportType, format: ExportFormat) => Promise; dateRange: DateRangePickerValue; - selectedFilters?: string[]; + selectedFilters?: readonly string[]; customTitle?: string; teams?: Team[]; } diff --git a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx index e27aab717ef..b56acfca86b 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx @@ -4,7 +4,7 @@ import type { SearchSelectOption } from "@/components/shared/SearchSelect"; import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; interface TeamMultiSelectProps { - value?: string[]; + value?: readonly string[]; onChange?: (value: string[]) => void; disabled?: boolean; organizationId?: string | null; diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx index 3757720c699..bad4356aafb 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx @@ -79,7 +79,7 @@ describe("MCPToolPermissions", () => { const expectedToolPermissions = { [mockServerId]: ["read_wiki_contents", "ask_question"], }; - expect(mockOnChange).toHaveBeenCalledWith(expectedToolPermissions); + expect(mockOnChange).toHaveBeenCalledWith(expectedToolPermissions, "edit"); // Verify API calls // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock @@ -135,9 +135,12 @@ describe("MCPToolPermissions", () => { // Selecting every displayed tool writes the wildcard, which also covers tools the // server adds later. - expect(mockOnChange).toHaveBeenCalledWith({ - [mockServerId]: ["*"], - }); + expect(mockOnChange).toHaveBeenCalledWith( + { + [mockServerId]: ["*"], + }, + "edit", + ); }); it("should deselect all tools when Deselect All button is clicked", async () => { @@ -186,9 +189,12 @@ describe("MCPToolPermissions", () => { await userEvent.click(deselectAllButton); // Verify onChange was called with no tools selected - expect(mockOnChange).toHaveBeenCalledWith({ - [mockServerId]: [], - }); + expect(mockOnChange).toHaveBeenCalledWith( + { + [mockServerId]: [], + }, + "edit", + ); }); describe("wildcard all-tools grant", () => { @@ -238,7 +244,7 @@ describe("MCPToolPermissions", () => { expect(await screen.findByText("read_wiki_structure")).toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: "Select All" })); - expect(mockOnChange).toHaveBeenCalledWith({ [wildcardServerId]: ["*"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [wildcardServerId]: ["*"] }, "edit"); }); it("converts back to an enumerated list when one tool is unchecked from a wildcard grant", async () => { @@ -256,9 +262,12 @@ describe("MCPToolPermissions", () => { await userEvent.click(screen.getByText("Flat List")); await userEvent.click(screen.getByRole("checkbox", { name: "ask_question" })); - expect(mockOnChange).toHaveBeenCalledWith({ - [wildcardServerId]: ["read_wiki_structure", "read_wiki_contents"], - }); + expect(mockOnChange).toHaveBeenCalledWith( + { + [wildcardServerId]: ["read_wiki_structure", "read_wiki_contents"], + }, + "edit", + ); }); }); @@ -320,7 +329,7 @@ describe("MCPToolPermissions", () => { expect(deleteIssue).toBeChecked(); await userEvent.click(listIssues); - expect(mockOnChange).toHaveBeenCalledWith({ [groupServer.server_id]: ["delete_issue"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [groupServer.server_id]: ["delete_issue"] }, "edit"); }); it("marks an access-group server as inherited and leaves a directly selected one unmarked", async () => { @@ -469,7 +478,7 @@ describe("MCPToolPermissions", () => { // First checkbox is the header toggle of the group holding list_issues. await userEvent.click(screen.getAllByRole("checkbox")[0]); - expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["list_issues"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["list_issues"] }, "edit"); }); // Copying the toolset's tools into the entry would outlive the toolset, so a write keeps only @@ -502,7 +511,7 @@ describe("MCPToolPermissions", () => { // A toolset-sourced server never writes the wildcard: that would create a standing direct // grant outliving the toolset. The write keeps only the tools this level grants itself. - expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["delete_issue"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["delete_issue"] }, "edit"); }); // The default narrows an unrestricted server; against a toolset-restricted one it would widen @@ -597,7 +606,7 @@ describe("MCPToolPermissions", () => { expect(await screen.findByText("Payments")).toBeInTheDocument(); expect(screen.queryByText("srv-collide")).not.toBeInTheDocument(); await waitFor(() => { - expect(mockOnChange).toHaveBeenCalledWith({ "srv-collide": ["list_issues"] }); + expect(mockOnChange).toHaveBeenCalledWith({ "srv-collide": ["list_issues"] }, "default"); }); expect(mockOnChange.mock.calls.every(([written]) => !Object.hasOwn(written, "srv-twin"))).toBe(true); expect(networking.listMCPTools).not.toHaveBeenCalledWith(mockAccessToken, "srv-twin"); @@ -640,7 +649,7 @@ describe("MCPToolPermissions", () => { ); await waitFor(() => { - expect(mockOnChange).toHaveBeenCalledWith({ [directServer.server_id]: ["list_issues"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [directServer.server_id]: ["list_issues"] }, "default"); }); }); @@ -804,7 +813,7 @@ describe("MCPToolPermissions", () => { expect(await screen.findByText("list_issues")).toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: "Deselect All" })); - expect(mockOnChange).toHaveBeenCalledWith({ github_mcp: [] }); + expect(mockOnChange).toHaveBeenCalledWith({ github_mcp: [] }, "edit"); }); }); @@ -984,7 +993,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(readGroupToggle); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["delete_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["delete_document"] }, "edit"); }); it("adds the rest of a partially-allowed risk group when its mixed toggle is clicked", async () => { @@ -996,7 +1005,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(readGroupToggle); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }, "edit"); }); it("toggles a single tool exactly once when its checkbox is clicked inside the clickable row", async () => { @@ -1006,7 +1015,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(await screen.findByRole("checkbox", { name: "delete_document" })); expect(mockOnChange).toHaveBeenCalledTimes(1); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }, "edit"); }); it("toggles a single tool when the row around its checkbox is clicked", async () => { @@ -1016,7 +1025,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(await screen.findByText("Destroy a document")); expect(mockOnChange).toHaveBeenCalledTimes(1); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }, "edit"); }); it("re-renders each checkbox from the permissions it emitted", async () => { @@ -1103,7 +1112,7 @@ describe("MCPToolPermissions", () => { />, ); - await waitFor(() => expect(onChange).toHaveBeenCalledWith({ "srv-files": ["read_file"] })); + await waitFor(() => expect(onChange).toHaveBeenCalledWith({ "srv-files": ["read_file"] }, "default")); }); }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index 8067936e715..b12b31601b8 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -20,13 +20,15 @@ import { resolveEffectiveMcpServers, } from "./effectiveMcpServers"; +export type ToolPermissionsChangeOrigin = "default" | "edit"; + interface MCPToolPermissionsProps { accessToken: string; selectedServers: readonly string[]; selectedAccessGroups?: readonly string[]; selectedToolsets?: readonly string[]; toolPermissions: Record; - onChange: (toolPermissions: Record) => void; + onChange: (toolPermissions: Record, origin: ToolPermissionsChangeOrigin) => void; disabled?: boolean; preserveUnrestrictedServers?: readonly string[]; } @@ -141,7 +143,10 @@ const MCPToolPermissions: React.FC = ({ const nonDeleteTools = fetchedTools .filter((t) => classifyToolOp(t.name, t.description || "") !== "delete") .map((t) => t.name); - onChange(applyToolPermissionWrite({ toolPermissions: latestPermissions, entry, allowed: nonDeleteTools })); + onChange( + applyToolPermissionWrite({ toolPermissions: latestPermissions, entry, allowed: nonDeleteTools }), + "default", + ); } } } catch (err) { @@ -175,7 +180,7 @@ const MCPToolPermissions: React.FC = ({ entry.source.kind !== "toolset" && names.length > 0 && names.every((n) => allowed.includes(n)) ? [MCP_ALL_TOOLS_WILDCARD] : allowed; - onChange(applyToolPermissionWrite({ toolPermissions, entry, allowed: next })); + onChange(applyToolPermissionWrite({ toolPermissions, entry, allowed: next }), "edit"); }; const handleSelectAll = (entry: EffectiveMcpServer) => { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 69623d05232..0f7d76a7156 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1465,7 +1465,7 @@ export const teamSpendByUserCall = async ( accessToken: string, startTime: Date, endTime: Date, - teamIds: string[], + teamIds: readonly string[], ): Promise => apiClient.get(`/team/spend/by_user`, { accessToken, diff --git a/ui/litellm-dashboard/src/components/per_user_usage.tsx b/ui/litellm-dashboard/src/components/per_user_usage.tsx index f600cd6c45d..a7c6a5cf504 100644 --- a/ui/litellm-dashboard/src/components/per_user_usage.tsx +++ b/ui/litellm-dashboard/src/components/per_user_usage.tsx @@ -26,7 +26,7 @@ interface PerUserAnalyticsResponse { interface PerUserUsageProps { accessToken: string | null; - selectedTags: string[]; + selectedTags: readonly string[]; formatAbbreviatedNumber: (value: number, decimalPlaces?: number) => string; } @@ -42,7 +42,7 @@ const PerUserUsage: React.FC = ({ accessToken, selectedTags, }); const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: 50 }); - const [pagedTags, setPagedTags] = useState(selectedTags); + const [pagedTags, setPagedTags] = useState(selectedTags); if (pagedTags !== selectedTags) { setPagedTags(selectedTags); @@ -57,7 +57,7 @@ const PerUserUsage: React.FC = ({ accessToken, selectedTags, accessToken, pagination.pageIndex + 1, pagination.pageSize, - pagedTags.length > 0 ? pagedTags : undefined, + pagedTags.length > 0 ? [...pagedTags] : undefined, ) .then((response) => { if (stale) return; diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 4d0d590521f..0ebd2cdff17 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -214,6 +214,7 @@ describe("provider_info_helpers", () => { Providers.LEMONADE, Providers.LLAMAFILE, Providers.MARITALK, + Providers.Nace, Providers.NLP_CLOUD, Providers.NSCALE, Providers.OVHCLOUD, @@ -517,6 +518,16 @@ describe("provider_info_helpers", () => { expect(getProviderModels("Sail" as Providers, modelMap)).toEqual(["sail/openai/gpt-oss-120b"]); }); + it("lists Nace decision models and resolves the provider name and model placeholder", () => { + const modelMap = { + "nace/drex-latest": { litellm_provider: "nace", mode: "evaluation" }, + "typesafe/jev-latest": { litellm_provider: "typesafe", mode: "evaluation" }, + }; + expect(getProviderModels(Providers.Nace, modelMap)).toEqual(["nace/drex-latest"]); + expect(getProviderLogoAndName("nace").displayName).toBe("Nace"); + expect(getPlaceholder("Nace")).toBe("nace/drex-latest"); + }); + it("should list only typesafe decision models for the 'TypeSafe' provider key, not the OpenRouter-hosted one", () => { const modelMap = { "typesafe/jev-latest": { litellm_provider: "typesafe", mode: "evaluation" }, diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 3925118a881..5a163bce3b6 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -181,6 +181,7 @@ export enum Providers { TOPAZ = "Topaz", Triton = "Triton", TypeSafe = "TypeSafe", + Nace = "Nace", V0 = "V0", VERCEL_AI_GATEWAY = "Vercel Ai Gateway", Vertex_AI = "Vertex AI (Anthropic, Gemini, etc.)", @@ -307,6 +308,7 @@ export const provider_map: Record = { TOPAZ: "topaz", Triton: "triton", TypeSafe: "typesafe", + Nace: "nace", V0: "v0", VERCEL_AI_GATEWAY: "vercel_ai_gateway", Vertex_AI: "vertex_ai", @@ -485,6 +487,7 @@ const providerPlaceholderMap: Partial> = { [Providers.StrandsDecider]: "strands_decider/strands-decider-2B-hobson-v19", [Providers.Tencent]: "tencent/deepseek-v4-pro", [Providers.TypeSafe]: "typesafe/jev-latest", + [Providers.Nace]: "nace/drex-latest", [Providers.Vertex_AI]: "gemini-pro", [Providers.VolcEngine]: "volcengine/", [Providers.Voyage]: "voyage/", diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx index 502078e03a9..9c96f60dee2 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx @@ -22,7 +22,7 @@ import { usePaginatedCombobox } from "./usePaginatedCombobox"; interface PaginatedMultiSelectProps { options: SearchSelectOption[]; - value?: string[]; + value?: readonly string[]; onValueChange: (value: string[]) => void; onSearchChange: (query: string) => void; onLoadMore: () => void; diff --git a/ui/litellm-dashboard/src/components/shared/advanced_date_picker.test.tsx b/ui/litellm-dashboard/src/components/shared/advanced_date_picker.test.tsx index 78e258d1a69..104eed54b2f 100644 --- a/ui/litellm-dashboard/src/components/shared/advanced_date_picker.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/advanced_date_picker.test.tsx @@ -1,5 +1,6 @@ import { render, screen, fireEvent, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import moment from "moment"; import { describe, it, expect, vi, beforeEach, beforeAll } from "vitest"; import AdvancedDatePicker from "./advanced_date_picker"; @@ -166,6 +167,82 @@ describe("AdvancedDatePicker", () => { }); }); + it("passes the clicked preset label when applying", async () => { + const { container } = render(); + + openDropdown(container); + fireEvent.click(screen.getByText("Today")); + fireEvent.click(screen.getByRole("button", { name: "Apply" })); + + expect(mockOnValueChange).toHaveBeenNthCalledWith(1, expect.anything(), "today"); + await waitFor(() => { + expect(mockOnValueChange).toHaveBeenNthCalledWith(2, expect.anything(), "today"); + }); + }); + + it("passes a null preset label when typed dates match today", async () => { + const { container } = render(); + const today = moment().format("YYYY-MM-DD"); + + openDropdown(container); + fireEvent.change(screen.getByDisplayValue("2025-01-01"), { target: { value: today } }); + fireEvent.change(screen.getByDisplayValue("2025-01-31"), { target: { value: today } }); + fireEvent.click(screen.getByRole("button", { name: "Apply" })); + + expect(mockOnValueChange).toHaveBeenNthCalledWith(1, expect.anything(), null); + await waitFor(() => { + expect(mockOnValueChange).toHaveBeenNthCalledWith(2, expect.anything(), null); + }); + }); + + it("resets the preset label when the parent switches to a custom range", async () => { + const { container, rerender } = render( + , + ); + + openDropdown(container); + fireEvent.click(screen.getByText("Today")); + fireEvent.click(screen.getByRole("button", { name: "Apply" })); + + expect(mockOnValueChange).toHaveBeenNthCalledWith(1, expect.anything(), "today"); + await waitFor(() => { + expect(mockOnValueChange).toHaveBeenNthCalledWith(2, expect.anything(), "today"); + }); + + mockOnValueChange.mockClear(); + const customValue = { + from: new Date("2025-02-01T00:00:00.000Z"), + to: new Date("2025-02-10T23:59:59.999Z"), + }; + rerender(); + + openDropdown(container); + fireEvent.click(screen.getByRole("button", { name: "Apply" })); + + expect(mockOnValueChange).toHaveBeenNthCalledWith(1, expect.anything(), null); + await waitFor(() => { + expect(mockOnValueChange).toHaveBeenNthCalledWith(2, expect.anything(), null); + }); + }); + + it("preserves the parent preset label when applying an unchanged range", async () => { + const value = { + from: moment().subtract(30, "days").startOf("day").toDate(), + to: moment().endOf("day").toDate(), + }; + const { container } = render( + , + ); + + openDropdown(container); + fireEvent.click(screen.getByRole("button", { name: "Apply" })); + + expect(mockOnValueChange).toHaveBeenNthCalledWith(1, expect.anything(), "30d"); + await waitFor(() => { + expect(mockOnValueChange).toHaveBeenNthCalledWith(2, expect.anything(), "30d"); + }); + }); + it("should select relative time option", () => { const { container } = render(); diff --git a/ui/litellm-dashboard/src/components/shared/advanced_date_picker.tsx b/ui/litellm-dashboard/src/components/shared/advanced_date_picker.tsx index 190e3de4a82..513db07409c 100644 --- a/ui/litellm-dashboard/src/components/shared/advanced_date_picker.tsx +++ b/ui/litellm-dashboard/src/components/shared/advanced_date_picker.tsx @@ -7,20 +7,21 @@ import React, { useCallback, useEffect, useRef, useState } from "react"; interface AdvancedDatePickerProps { value: DateRangePickerValue; - onValueChange: (value: DateRangePickerValue) => void; + onValueChange: (value: DateRangePickerValue, presetShortLabel: string | null) => void; + presetShortLabel?: string | null; label?: string; className?: string; showTimeRange?: boolean; align?: "left" | "right"; } -interface RelativeTimeOption { +export interface RelativeTimeOption { label: string; shortLabel: string; getValue: () => { from: Date; to: Date }; } -const relativeTimeOptions: RelativeTimeOption[] = [ +export const relativeTimeOptions: readonly RelativeTimeOption[] = [ { label: "Today", shortLabel: "today", @@ -63,12 +64,22 @@ const relativeTimeOptions: RelativeTimeOption[] = [ }, ]; +export function matchRelativeTimeOption(value: DateRangePickerValue): RelativeTimeOption | undefined { + const { from, to } = value; + if (!from || !to) return undefined; + return relativeTimeOptions.find((option) => { + const optionRange = option.getValue(); + return moment(from).isSame(optionRange.from, "day") && moment(to).isSame(optionRange.to, "day"); + }); +} + /** * Advanced Date Range Picker with dropdown, relative times, and custom inputs */ const AdvancedDatePicker: React.FC = ({ value, onValueChange, + presetShortLabel = null, label = "Select Time Range", className, showTimeRange = true, @@ -77,6 +88,7 @@ const AdvancedDatePicker: React.FC = ({ const [isOpen, setIsOpen] = useState(false); const [tempValue, setTempValue] = useState(value); const [selectedOption, setSelectedOption] = useState(null); + const [pickedPreset, setPickedPreset] = useState(null); // Custom date inputs only - removed time inputs const [startDate, setStartDate] = useState(""); @@ -84,24 +96,10 @@ const AdvancedDatePicker: React.FC = ({ const dropdownRef = useRef(null); - // Function to check if current value matches a relative time option - const getMatchingOption = useCallback((currentValue: DateRangePickerValue): string | null => { - if (!currentValue.from || !currentValue.to) return null; - - for (const option of relativeTimeOptions) { - const optionRange = option.getValue(); - - // Compare dates with some tolerance (to account for time differences) - const fromMatches = moment(currentValue.from).isSame(moment(optionRange.from), "day"); - const toMatches = moment(currentValue.to).isSame(moment(optionRange.to), "day"); - - if (fromMatches && toMatches) { - return option.shortLabel; - } - } - - return null; - }, []); + const getMatchingOption = useCallback( + (currentValue: DateRangePickerValue): string | null => matchRelativeTimeOption(currentValue)?.shortLabel ?? null, + [], + ); // Update selected option when value changes useEffect(() => { @@ -140,7 +138,8 @@ const AdvancedDatePicker: React.FC = ({ setEndDate(moment(value.to).format("YYYY-MM-DD")); } setTempValue(value); - }, [value]); + setPickedPreset(presetShortLabel ?? null); + }, [value, presetShortLabel]); // Close dropdown when clicking outside useEffect(() => { @@ -206,6 +205,7 @@ const AdvancedDatePicker: React.FC = ({ // Update local state to reflect the selection (don't apply immediately) setTempValue(newValue); setSelectedOption(option.shortLabel); + setPickedPreset(option.shortLabel); // Update the form inputs to reflect the selection setStartDate(moment(from).format("YYYY-MM-DD")); @@ -243,13 +243,13 @@ const AdvancedDatePicker: React.FC = ({ const handleApply = () => { if (tempValue.from && tempValue.to && validation.isValid) { // First call with immediate value for UI responsiveness - onValueChange(tempValue); + onValueChange(tempValue, pickedPreset); // Then do the same background adjustment logic as the original component requestIdleCallback( () => { const adjustedValue = adjustDateRange(tempValue); - onValueChange(adjustedValue); + onValueChange(adjustedValue, pickedPreset); }, { timeout: 100 }, ); @@ -273,6 +273,7 @@ const AdvancedDatePicker: React.FC = ({ // Reset selected option const matchingOption = getMatchingOption(value); setSelectedOption(matchingOption); + setPickedPreset(presetShortLabel ?? null); setIsOpen(false); }; @@ -367,7 +368,10 @@ const AdvancedDatePicker: React.FC = ({ setStartDate(e.target.value)} + onChange={(e) => { + setPickedPreset(null); + setStartDate(e.target.value); + }} className={`w-65 px-3 py-2 text-sm border rounded-md cursor-pointer hover:border-ring focus:border-info focus:ring-1 focus:ring-ring ${ !validation.isValid ? "border-destructive/30 focus:border-destructive focus:ring-red-200" @@ -382,7 +386,10 @@ const AdvancedDatePicker: React.FC = ({ setEndDate(e.target.value)} + onChange={(e) => { + setPickedPreset(null); + setEndDate(e.target.value); + }} className={`w-65 px-3 py-2 text-sm border rounded-md cursor-pointer hover:border-ring focus:border-info focus:ring-1 focus:ring-ring ${ !validation.isValid ? "border-destructive/30 focus:border-destructive focus:ring-red-200" diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index 35f09011ab9..ffbdd7ae086 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -1,7 +1,8 @@ import { z } from "zod"; -import { UseFormSetValue } from "react-hook-form"; +import { UseFormReturn, UseFormSetValue } from "react-hook-form"; import { KeyResponse } from "../key_team_helpers/key_list"; +import { ToolPermissionsChangeOrigin } from "../mcp_server_management/MCPToolPermissions"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; import { estimateChecks, estimateFields } from "./estimatedOutputTokens"; @@ -222,6 +223,13 @@ export const toSubmittedValues = ( }); export const clearMcpGrants = (setValue: UseFormSetValue) => { - setValue("mcp_servers_and_groups", { servers: [], accessGroups: [], toolsets: [] }); - setValue("mcp_tool_permissions", {}); + setValue("mcp_servers_and_groups", { servers: [], accessGroups: [], toolsets: [] }, { shouldDirty: true }); + setValue("mcp_tool_permissions", {}, { shouldDirty: true }); }; + +export const writeMcpToolPermissions = + (form: Pick, "getFieldState" | "setValue">) => + (toolPermissions: Record, origin: ToolPermissionsChangeOrigin) => + form.setValue("mcp_tool_permissions", toolPermissions, { + shouldDirty: origin === "edit" || form.getFieldState("mcp_servers_and_groups").isDirty, + }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index ead17fa0c95..4dd98b31781 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -10,6 +10,7 @@ import { getPoliciesList, getUiSettings, getPromptsList, + listMCPTools, modelAvailableCall, vectorStoreListCall, } from "../networking"; @@ -1011,6 +1012,179 @@ describe("KeyEditView", () => { }); }); + it("should preserve a budget draft when the same key refreshes before saving", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: MOCK_KEY_DATA, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + view.rerender( + , + ); + await chooseSelectOption(userEvent, screen.getByLabelText("Reset Budget"), "monthly"); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + const expected = { + max_budget: "12.5", + budget_duration: "30d", + rpm_limit: 75, + models: ["gpt-4"], + metadata: JSON.stringify({ note: "new" }, null, 2), + vector_stores: ["store-new"], + }; + await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expected), expect.any(Array))); + const [, dirtyFields] = onSubmit.mock.calls[0]; + expect(dirtyFields).toContain("max_budget"); + expect(dirtyFields).not.toContain("rpm_limit"); + }); + + const mcpServer = (server_id: string) => ({ + server_id, + server_name: server_id, + url: `https://${server_id}.example`, + transport: "http", + auth_type: "none", + created_at: "", + created_by: "", + updated_at: "", + updated_by: "", + }); + const mcpTools = (names: string[]) => ({ + tools: names.map((name) => ({ name })), + error: null, + message: null, + stack_trace: null, + }); + + it("should not mark MCP tool defaults as an edit when saving only a budget", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue([mcpServer("srv-a")]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools(["read_page", "delete_page"])); + testQueryClient.clear(); + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + , + ); + await waitFor(() => expect(listMCPTools).toHaveBeenCalledWith("test-token", "srv-a")); + + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalled()); + const [values, dirtyFields] = onSubmit.mock.calls[0]; + expect(values.mcp_tool_permissions).toHaveProperty("srv-a"); + expect(dirtyFields).toEqual(["max_budget"]); + vi.mocked(fetchMCPServers).mockResolvedValue([]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools([])); + }); + + it("should keep default tool limits for a newly added MCP server when the same key refreshes", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue([mcpServer("mcp-1")]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools(["read-page", "delete-page"])); + testQueryClient.clear(); + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: MOCK_KEY_DATA, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + + await userEvent.click(await screen.findByTestId("mcp-server-selector")); + await waitFor(() => expect(listMCPTools).toHaveBeenCalledWith("test-token", "mcp-1")); + view.rerender(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalled()); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ + rpm_limit: 75, + mcp_tool_permissions: { "mcp-1": ["read-page"] }, + }); + vi.mocked(fetchMCPServers).mockResolvedValue([]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools([])); + }); + + it("should start a fresh budget draft when switching between keys identified only by token_id", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: { ...MOCK_KEY_DATA, token: "", token_id: "key-one" }, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + view.rerender( + , + ); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ max_budget: 25 }), expect.any(Array)), + ); + }); + + it("should start a fresh budget draft when switching to another key", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: MOCK_KEY_DATA, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + view.rerender( + , + ); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ max_budget: 25 }), expect.any(Array)), + ); + }); + it("should persist a canonical budget_duration value, not a word-form the backend cannot parse", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); renderWithProviders( @@ -1615,6 +1789,46 @@ describe("KeyEditView", () => { }); }); + it("keeps the team change's cleared MCP grants when the same key refreshes before saving", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const keyData = { + ...MOCK_KEY_DATA, + organization_id: "org-1", + team_id: "group-maple", + object_permission: { mcp_servers: ["srv-a"], mcp_tool_permissions: { "srv-a": ["read"] } }, + }; + const renderEditor = (data: KeyResponse) => ( + {}} + onSubmit={onSubmit} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + /> + ); + const view = renderWithProviders(renderEditor(keyData)); + + await waitFor(() => { + expect(screen.getAllByRole("button", { name: "Clear" })).toHaveLength(2); + }); + await userEvent.click(screen.getAllByRole("button", { name: "Clear" })[1]); + await waitFor(() => { + expect(screen.getByTestId("mcp-server-selector")).toHaveAttribute("data-team-id", ""); + }); + view.rerender(renderEditor({ ...keyData, rpm_limit: 75 })); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: {}, + }); + }); + it("clears the organization and its dependent team in the update payload", async () => { const onSubmit = vi.fn().mockResolvedValue(undefined); renderWithProviders( @@ -1649,7 +1863,7 @@ describe("KeyEditView", () => { }); }); - it("should save an explicit project detach while keeping parents locked until the saved key changes", async () => { + it("should save an explicit project detach while keeping parents locked until reopening the editor", async () => { vi.mocked(getUiSettings).mockResolvedValue({ values: { enable_projects_ui: true } }); const onSubmit = vi.fn().mockResolvedValue(undefined); const onCancel = vi.fn(); @@ -1679,17 +1893,23 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: "Cancel" })); expect(onCancel).toHaveBeenCalledOnce(); expect(onSubmit).not.toHaveBeenCalled(); + // KeyInfoView unmounts the editor on cancel and after a successful save. + view.rerender(<>); view.rerender(renderEditor({ ...key })); await userEvent.click(await screen.findByRole("button", { name: "Detach from project" })); + view.rerender(renderEditor({ ...key, spend: 1 })); + expect(screen.getByRole("button", { name: "Keep project" })).toBeVisible(); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); const expectedDetach = { project_id: null, organization_id: "org-1", team_id: "group-maple", models: key.models }; await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expectedDetach), expect.any(Array)), ); expect(screen.getByRole("combobox", { name: "Team" })).toBeDisabled(); + view.rerender(<>); view.rerender(renderEditor({ ...key, project_id: null })); expect(screen.getByRole("combobox", { name: "Team" })).toBeEnabled(); expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); + view.rerender(<>); view.rerender(renderEditor(key, "Internal User")); expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); view.rerender(renderEditor(key, "Org Admin")); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 14724fcf6ef..44e9af0d316 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -46,6 +46,7 @@ import { toKeyEditFormValues, toSubmittedValues, clearMcpGrants, + writeMcpToolPermissions, } from "./keyEditFormValues"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { END_USER_BUDGET_HINT, EndUserBudgetSelect } from "../key_team_helpers/EndUserBudgetSelect"; @@ -89,7 +90,11 @@ interface KeyEditViewProps { premiumUser?: boolean; } -export function KeyEditView({ +const sessionKey = (keyData: KeyResponse) => keyData.token || keyData.token_id; + +export const KeyEditView = (props: KeyEditViewProps) => ; + +function KeyEditSession({ keyData, onCancel, onSubmit, @@ -105,7 +110,9 @@ export function KeyEditView({ const canEditEstimates = userRole != null && isProxyAdminRole(userRole); const estimateTooltip = estimateTooltips(canEditEstimates); const form = useZodForm(keyEditFormSchema, { - defaultValues: toKeyEditFormValues(keyData), + // Refresh untouched fields without erasing the user's in-progress edits. + values: toKeyEditFormValues(keyData), + resetOptions: { keepDirtyValues: true }, }); const [promptsList, setPromptsList] = useState([]); const [tagsList, setTagsList] = useState>({}); @@ -187,22 +194,16 @@ export function KeyEditView({ // Sync disabled callbacks with form when component mounts useEffect(() => { - form.setValue("disabled_callbacks", disabledCallbacks); + form.setValue("disabled_callbacks", disabledCallbacks, { shouldDirty: true }); }, [form, disabledCallbacks]); - useEffect(() => { - form.reset(toKeyEditFormValues(keyData)); - }, [keyData, form]); - // Sync auto-rotation state with form values useEffect(() => { - form.setValue("auto_rotate", autoRotationEnabled); + form.setValue("auto_rotate", autoRotationEnabled, { shouldDirty: true }); }, [autoRotationEnabled, form]); useEffect(() => { - if (rotationInterval) { - form.setValue("rotation_interval", rotationInterval); - } + if (rotationInterval) form.setValue("rotation_interval", rotationInterval, { shouldDirty: true }); }, [rotationInterval, form]); // Fetch tags for selector @@ -327,7 +328,7 @@ export function KeyEditView({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); setSelectedOrganizationId(orgId); - form.setValue("team_id", null); + form.setValue("team_id", null, { shouldDirty: true }); }; const handleTeamChange = (setField: (value: string | null) => void, teamId: string | null) => { @@ -336,12 +337,12 @@ export function KeyEditView({ const organizationId = teams?.find((t) => t.team_id === teamId)?.organization_id; if (!organizationId) return; setSelectedOrganizationId(organizationId); - form.setValue("organization_id", organizationId); + form.setValue("organization_id", organizationId, { shouldDirty: true }); }; const handleDisabledCallbacksChange = (internalValues: string[]) => { setDisabledCallbacks(mapInternalToDisplayNames(internalValues)); - form.setValue("disabled_callbacks", internalValues); + form.setValue("disabled_callbacks", internalValues, { shouldDirty: true }); }; const modelOptions = [ @@ -417,14 +418,14 @@ export function KeyEditView({ onChange={(value) => { switch (value) { case "default": - form.setValue("allowed_routes", ""); + form.setValue("allowed_routes", "", { shouldDirty: true }); break; case "llm_api": - form.setValue("allowed_routes", "llm_api_routes"); + form.setValue("allowed_routes", "llm_api_routes", { shouldDirty: true }); break; case "management": - form.setValue("allowed_routes", "management_routes"); - form.setValue("models", []); + form.setValue("allowed_routes", "management_routes", { shouldDirty: true }); + form.setValue("models", [], { shouldDirty: true }); break; } }} @@ -775,7 +776,7 @@ export function KeyEditView({ selectedAccessGroups={mcpSelection?.accessGroups || []} selectedToolsets={mcpSelection?.toolsets || []} toolPermissions={(mcpToolPermissions as Record | undefined) || {}} - onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms)} + onChange={writeMcpToolPermissions(form)} />
@@ -827,7 +828,9 @@ export function KeyEditView({ canDetach={canDetachProject} pending={detachProject} disabled={isKeySaving} - onToggle={() => form.setValue("project_id", detachProject ? keyData.project_id : null)} + onToggle={() => + form.setValue("project_id", detachProject ? keyData.project_id : null, { shouldDirty: true }) + } /> )} diff --git a/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx b/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx index 0ea84dea970..9d5edd7e2cc 100644 --- a/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx +++ b/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx @@ -1,4 +1,5 @@ import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { STACKED_USAGE_PALETTE } from "@/components/shared/charts"; import UserAgentActivity from "./user_agent_activity"; @@ -11,11 +12,7 @@ vi.mock("./networking", () => ({ tagWauCall: vi.fn(), tagMauCall: vi.fn(), tagDistinctCall: vi.fn(), -})); - -// Mock PerUserUsage component -vi.mock("./per_user_usage", () => ({ - default: () =>
Per User Usage
, + perUserAnalyticsCall: vi.fn(), })); describe("UserAgentActivity", () => { @@ -24,6 +21,7 @@ describe("UserAgentActivity", () => { const mockTagWauCall = vi.mocked(networking.tagWauCall); const mockTagMauCall = vi.mocked(networking.tagMauCall); const mockTagDistinctCall = vi.mocked(networking.tagDistinctCall); + const mockPerUserAnalyticsCall = vi.mocked(networking.perUserAnalyticsCall); const mockDistinctTagsData = { results: [{ tag: "User-Agent: Chrome/1.0" }, { tag: "User-Agent: Firefox/2.0" }, { tag: "User-Agent: Safari/3.0" }], @@ -102,12 +100,20 @@ describe("UserAgentActivity", () => { mockTagWauCall.mockClear(); mockTagMauCall.mockClear(); mockTagDistinctCall.mockClear(); + mockPerUserAnalyticsCall.mockClear(); mockTagDistinctCall.mockResolvedValue(mockDistinctTagsData); mockUserAgentSummaryCall.mockResolvedValue(mockSummaryData); mockTagDauCall.mockResolvedValue(mockDauData); mockTagWauCall.mockResolvedValue(mockWauData); mockTagMauCall.mockResolvedValue(mockMauData); + mockPerUserAnalyticsCall.mockResolvedValue({ + results: [], + total_count: 0, + page: 1, + page_size: 50, + total_pages: 0, + }); }); it("should render summary cards with user agent data", async () => { @@ -185,6 +191,70 @@ describe("UserAgentActivity", () => { expect(prompts).toBeGreaterThan(0); }); + it("uses and reports a controlled tag selection", async () => { + const onSelectedTagsChange = vi.fn(); + const selectedTags = ["User-Agent: Chrome/1.0"]; + + render( + , + ); + + await waitFor(() => { + expect(mockTagDauCall).toHaveBeenCalledWith("test-token", expect.any(Date), undefined, selectedTags); + }); + fireEvent.click(await screen.findByRole("button", { name: "Clear user agent filter" })); + + expect(onSelectedTagsChange).toHaveBeenCalledWith([]); + }); + + it("keeps per-user pagination when controlled tags stay unchanged on parent rerender", async () => { + const user = userEvent.setup(); + const selectedTags = ["User-Agent: Chrome/1.0"]; + const onSelectedTagsChange = vi.fn(); + const totalUsers = 120; + const pageOfUsers = (page: number, pageSize: number) => { + const start = (page - 1) * pageSize; + const count = Math.max(0, Math.min(pageSize, totalUsers - start)); + return Array.from({ length: count }, (_, index) => ({ + user_id: `user-${start + index + 1}`, + user_email: null, + user_agent: "Chrome/1.0", + successful_requests: 5, + failed_requests: 0, + total_requests: 5, + total_tokens: 100, + spend: 1, + })); + }; + mockPerUserAnalyticsCall.mockImplementation(async (_token, page = 1, pageSize = 50) => ({ + results: pageOfUsers(page, pageSize), + total_count: totalUsers, + page, + page_size: pageSize, + total_pages: Math.ceil(totalUsers / pageSize), + })); + const props = { ...defaultProps, selectedTags, onSelectedTagsChange }; + const { rerender } = render(); + + await waitFor(() => expect(mockTagDistinctCall).toHaveBeenCalled()); + await waitFor(() => expect(mockUserAgentSummaryCall).toHaveBeenCalled()); + await waitFor(() => expect(mockTagDauCall).toHaveBeenCalled()); + await waitFor(() => expect(mockTagWauCall).toHaveBeenCalled()); + await waitFor(() => expect(mockTagMauCall).toHaveBeenCalled()); + await user.click(screen.getByRole("radio", { name: "Per user" })); + expect(await screen.findByText("user-1")).toBeInTheDocument(); + await user.click(screen.getByTestId("pagination-next")); + expect(await screen.findByText("user-51")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 51-100 of 120"); + + const callsBeforeRerender = mockPerUserAnalyticsCall.mock.calls.length; + rerender(); + + expect(mockPerUserAnalyticsCall).toHaveBeenCalledTimes(callsBeforeRerender); + expect(screen.getByText("user-51")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 51-100 of 120"); + }); + // Walks up from the panel's heading to the nearest ancestor that owns a chart, so the // assertions do not depend on how many wrappers the tab library puts around a panel. const chartForTitle = (title: string): HTMLElement => { diff --git a/ui/litellm-dashboard/src/components/user_agent_activity.tsx b/ui/litellm-dashboard/src/components/user_agent_activity.tsx index ca318bfb4fe..37b320e2b5c 100644 --- a/ui/litellm-dashboard/src/components/user_agent_activity.tsx +++ b/ui/litellm-dashboard/src/components/user_agent_activity.tsx @@ -123,6 +123,8 @@ interface UserAgentActivityProps { onDateChange?: (value: DateRangePickerValue) => void; // Optional - not used anymore /** Tags to pre-select, e.g. an agent's user-agent tags when opened from the Overview's Top agents. */ initialTags?: readonly string[]; + selectedTags?: readonly string[]; + onSelectedTagsChange?: (tags: readonly string[]) => void; } const UserAgentActivity: React.FC = ({ @@ -131,6 +133,8 @@ const UserAgentActivity: React.FC = ({ dateValue, onDateChange, initialTags, + selectedTags: selectedTagsProp, + onSelectedTagsChange, }) => { const anchor = useComboboxAnchor(); // Maximum number of categories to show in charts to prevent color palette overflow @@ -146,7 +150,15 @@ const UserAgentActivity: React.FC = ({ // Tag filtering state const [availableTags, setAvailableTags] = useState([]); - const [selectedTags, setSelectedTags] = useState(() => [...(initialTags ?? [])]); + const [localSelectedTags, setLocalSelectedTags] = useState(() => [...(initialTags ?? [])]); + const selectedTags = selectedTagsProp ?? localSelectedTags; + const updateSelectedTags = (next: string[]) => { + if (onSelectedTagsChange) { + onSelectedTagsChange(next); + return; + } + setLocalSelectedTags(next); + }; const [tagsLoading, setTagsLoading] = useState(false); // Separate loading states for each endpoint @@ -184,7 +196,7 @@ const UserAgentActivity: React.FC = ({ accessToken, today, userAgentFilter || undefined, - selectedTags.length > 0 ? selectedTags : undefined, + selectedTags.length > 0 ? [...selectedTags] : undefined, ); setDauData(data); } catch (error) { @@ -203,7 +215,7 @@ const UserAgentActivity: React.FC = ({ accessToken, today, userAgentFilter || undefined, - selectedTags.length > 0 ? selectedTags : undefined, + selectedTags.length > 0 ? [...selectedTags] : undefined, ); setWauData(data); } catch (error) { @@ -222,7 +234,7 @@ const UserAgentActivity: React.FC = ({ accessToken, today, userAgentFilter || undefined, - selectedTags.length > 0 ? selectedTags : undefined, + selectedTags.length > 0 ? [...selectedTags] : undefined, ); setMauData(data); } catch (error) { @@ -241,7 +253,7 @@ const UserAgentActivity: React.FC = ({ accessToken, dateValue.from, dateValue.to, - selectedTags.length > 0 ? selectedTags : undefined, + selectedTags.length > 0 ? [...selectedTags] : undefined, ); setSummaryData(summary); } catch (error) { @@ -455,12 +467,7 @@ const UserAgentActivity: React.FC = ({ action={
- setSelectedTags(next)} - > + } className="min-w-0 flex-1" aria-busy={tagsLoading}> {(selected: string[]) => diff --git a/ui/litellm-dashboard/src/lib/toast.test.ts b/ui/litellm-dashboard/src/lib/toast.test.ts index 0afd0ec2cc4..b4a776f123a 100644 --- a/ui/litellm-dashboard/src/lib/toast.test.ts +++ b/ui/litellm-dashboard/src/lib/toast.test.ts @@ -37,6 +37,11 @@ describe("toast", () => { expect(sonner.success).toHaveBeenCalledWith("saved", { description: "Model x", duration: 1500 }); }); + it("passes an id through so repeated toasts collapse into one", () => { + toast.warning("no access", { id: "denied" }); + expect(sonner.warning).toHaveBeenCalledWith("no access", expect.objectContaining({ id: "denied" })); + }); + it("dismiss clears every toast", () => { toast.dismiss(); expect(sonner.dismiss).toHaveBeenCalledWith(); diff --git a/ui/litellm-dashboard/src/lib/toast.ts b/ui/litellm-dashboard/src/lib/toast.ts index a057d144e6f..a7a4ce45ad0 100644 --- a/ui/litellm-dashboard/src/lib/toast.ts +++ b/ui/litellm-dashboard/src/lib/toast.ts @@ -7,6 +7,7 @@ export type ToastKind = "success" | "info" | "warning" | "error"; export type ToastOptions = { readonly description?: ReactNode; readonly durationMs?: number; + readonly id?: string; }; type ErrorFacts = { @@ -127,6 +128,7 @@ const titleFor = ({ status, proxyType }: ErrorFacts): string => { const show = (kind: ToastKind, message: ReactNode, options?: ToastOptions): void => { sonner[kind](message, { + id: options?.id, description: options?.description, duration: options?.durationMs ?? DEFAULT_DURATION_MS[kind], }); diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index b3cd45a15f0..31e276914b0 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -26,12 +26,18 @@ export const testQueryClient = new QueryClient({ interface ProviderOptions { searchParams?: string | Record | URLSearchParams; onUrlUpdate?: OnUrlUpdateFunction; + resetUrlUpdateQueueOnMount?: boolean; } export const renderWithProviders = (ui: React.ReactElement, options?: RenderOptions & ProviderOptions) => { - const { searchParams, onUrlUpdate, ...renderOptions } = options ?? {}; + const { searchParams, onUrlUpdate, resetUrlUpdateQueueOnMount, ...renderOptions } = options ?? {}; const Providers: React.FC = ({ children }) => ( - + {children}